precommit fixes
Taylor committed
Oct 5, 2023 at 07:59 UTC
1237444ba7a6729c17799f8b646c420d72a6d8e9
47 files changed
+35
-359
backend/app/connectors/graylog/schema/monitoring.py
-3
@@ -1,8 +1,5 @@
1
-from typing import Any
2
-from typing import Dict
1
from typing import List
2
from typing import Optional
5
-from typing import Union
3
4
from pydantic import BaseModel
5
from pydantic import Field
backend/app/connectors/graylog/services/collector.py
-7
@@ -1,14 +1,7 @@
1
-from typing import Any
2
-from typing import Dict
1
from typing import List
4
-from typing import Optional
2
from typing import Tuple
6
-from typing import Union
3
8
-import requests
9
-import xmltodict
4
from loguru import logger
11
-from pydantic import Field
5
6
from app.connectors.graylog.schema.collector import ConfiguredInput
7
from app.connectors.graylog.schema.collector import ConfiguredInputsResponse
backend/app/connectors/graylog/services/events.py
-11
@@ -1,20 +1,9 @@
1
-from typing import Any
2
-from typing import Dict
3
-from typing import List
4
-from typing import Optional
5
-from typing import Tuple
6
-from typing import Union
7
-
8
-import requests
9
-import xmltodict
1
from loguru import logger
11
-from pydantic import Field
2
3
from app.connectors.graylog.schema.events import AlertEvent
4
from app.connectors.graylog.schema.events import AlertQuery
5
from app.connectors.graylog.schema.events import Alerts
6
from app.connectors.graylog.schema.events import Context
17
-from app.connectors.graylog.schema.events import Event
7
from app.connectors.graylog.schema.events import EventDefinition
8
from app.connectors.graylog.schema.events import GraylogAlertsResponse
9
from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse
backend/app/connectors/graylog/services/management.py
-11
@@ -1,14 +1,4 @@
1
-from typing import Any
2
-from typing import Dict
3
-from typing import List
4
-from typing import Optional
5
-from typing import Tuple
6
-from typing import Union
7
-
8
-import requests
9
-import xmltodict
1
from loguru import logger
11
-from pydantic import Field
2
3
from app.connectors.graylog.schema.management import DeletedIndexBody
4
from app.connectors.graylog.schema.management import DeletedIndexResponse
@@ -22,7 +12,6 @@ from app.connectors.graylog.schema.management import StopStreamBody
12
from app.connectors.graylog.schema.management import StopStreamResponse
13
from app.connectors.graylog.services.collector import get_index_names
14
from app.connectors.graylog.utils.universal import send_delete_request
25
-from app.connectors.graylog.utils.universal import send_get_request
15
from app.connectors.graylog.utils.universal import send_post_request
16
from app.connectors.graylog.utils.universal import send_put_request
17
backend/app/connectors/graylog/services/monitoring.py
-12
@@ -1,22 +1,10 @@
1
-from typing import Any
2
-from typing import Dict
3
-from typing import List
4
-from typing import Optional
5
-from typing import Tuple
6
-from typing import Union
7
-
8
-import requests
9
-import xmltodict
1
from loguru import logger
11
-from pydantic import Field
2
3
from app.connectors.graylog.schema.monitoring import GraylogMessages
4
from app.connectors.graylog.schema.monitoring import GraylogMessagesResponse
5
from app.connectors.graylog.schema.monitoring import GraylogMetricsResponse
6
from app.connectors.graylog.schema.monitoring import GraylogThroughputMetrics
7
from app.connectors.graylog.schema.monitoring import GraylogThroughputMetricsCollection
18
-from app.connectors.graylog.schema.monitoring import GraylogThroughputMetricsList
19
-from app.connectors.graylog.schema.monitoring import GraylogTotalMessages
8
from app.connectors.graylog.schema.monitoring import GraylogUncommittedJournalEntries
9
from app.connectors.graylog.utils.universal import send_get_request
10
backend/app/connectors/graylog/services/pipelines.py
-9
@@ -1,22 +1,13 @@
1
from typing import Any
2
from typing import Dict
3
-from typing import List
4
-from typing import Optional
5
-from typing import Tuple
6
-from typing import Union
3
8
-import requests
9
-import xmltodict
4
from loguru import logger
11
-from pydantic import Field
5
6
from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse
7
from app.connectors.graylog.schema.pipelines import Pipeline
8
from app.connectors.graylog.schema.pipelines import PipelineRule
9
from app.connectors.graylog.schema.pipelines import PipelineRulesResponse
17
-from app.connectors.graylog.schema.pipelines import Stage
10
from app.connectors.graylog.utils.universal import send_get_request
19
-from app.connectors.graylog.utils.universal import send_post_request
11
12
13
def get_pipelines() -> GraylogPipelinesResponse:
backend/app/connectors/graylog/services/streams.py
-9
@@ -1,19 +1,10 @@
1
-from typing import Any
2
-from typing import Dict
1
from typing import List
4
-from typing import Optional
5
-from typing import Tuple
6
-from typing import Union
2
8
-import requests
9
-import xmltodict
3
from loguru import logger
11
-from pydantic import Field
4
5
from app.connectors.graylog.schema.streams import GraylogStreamsResponse
6
from app.connectors.graylog.schema.streams import Stream
7
from app.connectors.graylog.utils.universal import send_get_request
16
-from app.connectors.graylog.utils.universal import send_post_request
8
9
10
def get_streams() -> GraylogStreamsResponse:
backend/app/connectors/graylog/utils/universal.py
-9
@@ -1,20 +1,11 @@
1
from typing import Any
2
from typing import Dict
3
-from typing import Generator
4
-from typing import List
3
from typing import Optional
6
-from typing import Type
4
5
import requests
9
-from elasticsearch7 import Elasticsearch
6
from loguru import logger
11
-from sqlmodel import Session
12
-from sqlmodel import select
7
14
-from app.connectors.models import Connectors
15
-from app.connectors.schema import ConnectorResponse
8
from app.connectors.utils import get_connector_info_from_db
17
-from app.db.db_session import engine
9
10
HEADERS = {"X-Requested-By": "CoPilot"}
11
backend/app/connectors/routes.py
-17
@@ -1,4 +1,3 @@
1
-from typing import List
1
from typing import Union
2
3
## Auth Things
@@ -6,33 +5,17 @@ from fastapi import APIRouter
5
from fastapi import Depends
6
from fastapi import File
7
from fastapi import HTTPException
9
-from fastapi import Request
10
-from fastapi import Security
8
from fastapi import UploadFile
12
-from fastapi import security
13
-from fastapi.responses import FileResponse
14
-from fastapi.responses import JSONResponse
15
-from fastapi.security import HTTPAuthorizationCredentials
9
from loguru import logger
17
-from starlette.status import HTTP_204_NO_CONTENT
10
from starlette.status import HTTP_401_UNAUTHORIZED
19
-from starlette.status import HTTP_404_NOT_FOUND
11
21
-from app.auth.models.users import User
22
-from app.auth.models.users import UserInput
23
-from app.auth.models.users import UserLogin
12
from app.auth.routes.auth import auth_handler
25
-from app.auth.schema.auth import UserLoginResponse
26
-from app.auth.schema.auth import UserResponse
27
-from app.auth.services.universal import find_user
28
-from app.auth.services.universal import select_all_users
13
from app.connectors.schema import ConnectorListResponse
14
from app.connectors.schema import ConnectorResponse
15
from app.connectors.schema import ConnectorsListResponse
16
from app.connectors.schema import UpdateConnector
17
from app.connectors.schema import VerifyConnectorResponse
18
from app.connectors.services import ConnectorServices
35
-from app.db.db_session import session
19
20
connector_router = APIRouter()
21
backend/app/connectors/schema.py
-2
@@ -1,6 +1,4 @@
1
from datetime import datetime
2
-from typing import Any
3
-from typing import Dict
2
from typing import List
3
from typing import Optional
4
backend/app/connectors/services.py
-2
@@ -1,8 +1,6 @@
1
import os
2
from contextlib import contextmanager
3
from datetime import datetime
4
-from typing import Any
5
-from typing import Dict
4
from typing import Generator
5
from typing import List
6
from typing import Optional
backend/app/connectors/shuffle/routes/workflows.py
+3
-10
@@ -1,24 +1,17 @@
1
-import json
2
-from typing import List
3
-
4
-import pydantic
1
from fastapi import APIRouter
6
-from fastapi import Depends
2
from fastapi import HTTPException
8
-from fastapi import Security
3
from loguru import logger
10
-from starlette.status import HTTP_401_UNAUTHORIZED
4
12
-# App specific imports
13
-from app.auth.routes.auth import auth_handler
5
from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel
6
from app.connectors.shuffle.schema.workflows import WorkflowExecutionResponseModel
16
-from app.connectors.shuffle.schema.workflows import WorkflowExecutionStatusResponseModel
7
from app.connectors.shuffle.schema.workflows import WorkflowsResponse
8
from app.connectors.shuffle.services.workflows import get_workflow_executions
9
from app.connectors.shuffle.services.workflows import get_workflows
10
from app.db.db_session import session
11
12
+# App specific imports
13
+
14
+
15
shuffle_workflows_router = APIRouter()
16
17
backend/app/connectors/shuffle/schema/workflows.py
-1
@@ -2,7 +2,6 @@ from typing import Any
2
from typing import Dict
3
from typing import List
4
from typing import Optional
5
-from typing import Union
5
6
from pydantic import BaseModel
7
from pydantic import Field
backend/app/connectors/shuffle/services/workflows.py
-12
@@ -1,18 +1,6 @@
1
-import json
2
-from typing import Any
3
-from typing import Dict
4
-from typing import List
5
-from typing import Optional
6
-from typing import Tuple
7
-from typing import Union
8
-
9
-import requests
10
-import xmltodict
1
from loguru import logger
12
-from pydantic import Field
2
3
from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel
15
-from app.connectors.shuffle.schema.workflows import WorkflowExecutionResponseModel
4
from app.connectors.shuffle.schema.workflows import WorkflowExecutionStatusResponseModel
5
from app.connectors.shuffle.schema.workflows import WorkflowsResponse
6
from app.connectors.shuffle.utils.universal import send_get_request
backend/app/connectors/shuffle/utils/universal.py
+9
-9
@@ -1,20 +1,11 @@
1
from typing import Any
2
from typing import Dict
3
-from typing import Generator
4
-from typing import List
3
from typing import Optional
6
-from typing import Type
4
5
import requests
9
-from elasticsearch7 import Elasticsearch
6
from loguru import logger
11
-from sqlmodel import Session
12
-from sqlmodel import select
7
14
-from app.connectors.models import Connectors
15
-from app.connectors.schema import ConnectorResponse
8
from app.connectors.utils import get_connector_info_from_db
17
-from app.db.db_session import engine
9
10
11
def verify_shuffle_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
@@ -120,6 +111,9 @@ def send_post_request(endpoint: str, data: Dict[str, Any] = None, connector_name
111
return {"success": False, "message": "No Graylog connector found in the database"}
112
113
try:
114
+ HEADERS = {
115
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
116
+ }
117
response = requests.post(
118
f"{attributes['connector_url']}{endpoint}",
119
headers=HEADERS,
@@ -163,6 +157,9 @@ def send_delete_request(endpoint: str, params: Optional[Dict[str, Any]] = None,
157
logger.error("No Graylog connector found in the database")
158
return None
159
try:
160
+ HEADERS = {
161
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
162
+ }
163
response = requests.delete(
164
f"{attributes['connector_url']}{endpoint}",
165
headers=HEADERS,
@@ -197,6 +194,9 @@ def send_put_request(endpoint: str, data: Optional[Dict[str, Any]] = None, conne
194
logger.error("No Graylog connector found in the database")
195
return None
196
try:
197
+ HEADERS = {
198
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
199
+ }
200
response = requests.put(
201
f"{attributes['connector_url']}{endpoint}",
202
headers=HEADERS,
backend/app/connectors/sublime/routes/alerts.py
-12
@@ -1,23 +1,11 @@
1
-import json
2
-from typing import Any
3
-from typing import List
4
-
5
-import pydantic
1
from fastapi import APIRouter
7
-from fastapi import Depends
8
-from fastapi import HTTPException
9
-from fastapi import Security
2
from loguru import logger
11
-from starlette.status import HTTP_401_UNAUTHORIZED
3
13
-# App specific imports
14
-from app.auth.routes.auth import auth_handler
4
from app.connectors.sublime.schema.alerts import AlertRequestBody
5
from app.connectors.sublime.schema.alerts import AlertResponseBody
6
from app.connectors.sublime.schema.alerts import SublimeAlertsResponse
7
from app.connectors.sublime.services.alerts import collect_alerts
8
from app.connectors.sublime.services.alerts import store_sublime_alert
20
-from app.db.db_session import session
9
10
sublime_alerts_router = APIRouter()
11
backend/app/connectors/sublime/utils/universal.py
-9
@@ -1,20 +1,11 @@
1
from typing import Any
2
from typing import Dict
3
-from typing import Generator
4
-from typing import List
3
from typing import Optional
6
-from typing import Type
4
5
import requests
9
-from elasticsearch7 import Elasticsearch
6
from loguru import logger
11
-from sqlmodel import Session
12
-from sqlmodel import select
7
14
-from app.connectors.models import Connectors
15
-from app.connectors.schema import ConnectorResponse
8
from app.connectors.utils import get_connector_info_from_db
17
-from app.db.db_session import engine
9
10
11
def verify_sublime_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
backend/app/connectors/utils.py
-7
@@ -1,14 +1,7 @@
1
-from contextlib import contextmanager
2
-from datetime import datetime
1
from typing import Any
2
from typing import Dict
5
-from typing import Generator
6
-from typing import List
7
-from typing import Optional
8
-from typing import Type
3
4
from loguru import logger
11
-from pydantic import BaseModel
5
from sqlmodel import Session
6
from sqlmodel import select
7
backend/app/connectors/velociraptor/routes/artifacts.py
+4
-9
@@ -1,18 +1,10 @@
1
-from datetime import timedelta
2
-from typing import Dict
1
from typing import List
4
-from typing import Optional
5
-from typing import Union
2
3
from fastapi import APIRouter
4
from fastapi import Depends
5
from fastapi import HTTPException
10
-from fastapi import Security
6
from loguru import logger
12
-from starlette.status import HTTP_401_UNAUTHORIZED
7
14
-# App specific imports
15
-from app.auth.routes.auth import auth_handler
8
from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
9
from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
10
from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
@@ -29,6 +21,9 @@ from app.connectors.velociraptor.services.artifacts import run_remote_command
21
from app.db.db_session import session
22
from app.db.universal_models import Agents
23
24
+# App specific imports
25
+
26
+
27
velociraptor_artifacts_router = APIRouter()
28
29
@@ -74,7 +69,7 @@ def get_velociraptor_id(hostname: str) -> str:
69
70
@velociraptor_artifacts_router.get("", response_model=ArtifactsResponse, description="Get all artifacts")
71
async def get_all_artifacts() -> ArtifactsResponse:
77
- logger.info(f"Fetching all artifacts")
72
+ logger.info("Fetching all artifacts")
73
return get_artifacts()
74
75
backend/app/connectors/velociraptor/services/artifacts.py
+1
-7
@@ -1,7 +1,3 @@
1
-from typing import Any
2
-from typing import Dict
3
-from typing import Optional
4
-
1
from fastapi import HTTPException
2
from loguru import logger
3
@@ -14,8 +10,6 @@ from app.connectors.velociraptor.schema.artifacts import QuarantineResponse
10
from app.connectors.velociraptor.schema.artifacts import RunCommandBody
11
from app.connectors.velociraptor.schema.artifacts import RunCommandResponse
12
from app.connectors.velociraptor.utils.universal import UniversalService
17
-from app.db.db_session import session
18
-from app.db.universal_models import Agents
13
14
universal_service = UniversalService()
15
@@ -66,7 +60,7 @@ def get_artifacts() -> ArtifactsResponse:
60
Returns:
61
ArtifactsResponse: A dictionary containing the artifacts.
62
"""
69
- logger.info(f"Fetching artifacts from Velociraptor")
63
+ logger.info("Fetching artifacts from Velociraptor")
64
query = create_query("SELECT name,description FROM artifact_definitions()")
65
all_artifacts = universal_service.execute_query(query)
66
if all_artifacts["success"]:
backend/app/connectors/velociraptor/utils/universal.py
-8
@@ -2,22 +2,14 @@ import json
2
from datetime import datetime
3
from typing import Any
4
from typing import Dict
5
-from typing import Optional
5
6
import grpc
8
-import pika
7
import pyvelociraptor
10
-import requests
8
from loguru import logger
9
from pyvelociraptor import api_pb2
10
from pyvelociraptor import api_pb2_grpc
14
-from sqlmodel import Session
15
-from sqlmodel import select
11
17
-from app.connectors.models import Connectors
18
-from app.connectors.schema import ConnectorResponse
12
from app.connectors.utils import get_connector_info_from_db
20
-from app.db.db_session import engine
13
14
15
def verify_velociraptor_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
backend/app/connectors/wazuh_indexer/routes/alerts.py
+3
-5
@@ -3,12 +3,8 @@ from typing import List
3
from fastapi import APIRouter
4
from fastapi import Depends
5
from fastapi import HTTPException
6
-from fastapi import Security
6
from loguru import logger
8
-from starlette.status import HTTP_401_UNAUTHORIZED
7
10
-# App specific imports
11
-from app.auth.routes.auth import auth_handler
8
from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse
9
from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHostResponse
10
from app.connectors.wazuh_indexer.schema.alerts import AlertsByRuleResponse
@@ -25,7 +21,9 @@ from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule_per_
21
from app.connectors.wazuh_indexer.services.alerts import get_host_alerts
22
from app.connectors.wazuh_indexer.services.alerts import get_index_alerts
23
from app.connectors.wazuh_indexer.utils.universal import collect_indices
28
-from app.db.db_session import session
24
+
25
+# App specific imports
26
+
27
28
wazuh_indexer_alerts_router = APIRouter()
29
backend/app/connectors/wazuh_indexer/schema/monitoring.py
-2
@@ -1,5 +1,3 @@
1
-from typing import Any
2
-from typing import Dict
1
from typing import List
2
from typing import Optional
3
from typing import Union
backend/app/connectors/wazuh_indexer/services/alerts.py
-15
@@ -1,21 +1,11 @@
1
-from typing import Any
1
from typing import Dict
2
from typing import List
3
from typing import Optional
4
from typing import Type
6
-from typing import Union
5
8
-import requests
9
-from elasticsearch7 import Elasticsearch
6
from fastapi import HTTPException
7
from loguru import logger
12
-from pydantic import BaseModel
13
-from sqlmodel import Session
14
-from sqlmodel import select
8
16
-from app.connectors.models import Connectors
17
-from app.connectors.schema import ConnectorResponse
18
-from app.connectors.utils import get_connector_info_from_db
9
from app.connectors.wazuh_indexer.schema.alerts import AlertsByHost
10
from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse
11
from app.connectors.wazuh_indexer.schema.alerts import AlertsByRule
@@ -29,14 +19,9 @@ from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchBody
19
from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchResponse
20
from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchBody
21
from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchResponse
32
-from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel
22
from app.connectors.wazuh_indexer.utils.universal import AlertsQueryBuilder
23
from app.connectors.wazuh_indexer.utils.universal import collect_indices
24
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
36
-from app.connectors.wazuh_indexer.utils.universal import format_indices_stats
37
-from app.connectors.wazuh_indexer.utils.universal import format_node_allocation
38
-from app.connectors.wazuh_indexer.utils.universal import format_shards
39
-from app.db.db_session import engine
25
26
# def collect_and_aggregate_alerts(field_name: str, search_body: AlertsSearchBody) -> Dict[str, int]:
27
# indices = collect_indices()
backend/app/connectors/wazuh_indexer/services/monitoring.py
-9
@@ -1,16 +1,8 @@
1
-from typing import Any
1
from typing import Dict
2
from typing import Union
3
5
-import requests
6
-from elasticsearch7 import Elasticsearch
4
from loguru import logger
8
-from sqlmodel import Session
9
-from sqlmodel import select
5
11
-from app.connectors.models import Connectors
12
-from app.connectors.schema import ConnectorResponse
13
-from app.connectors.utils import get_connector_info_from_db
6
from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealth
7
from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealthResponse
8
from app.connectors.wazuh_indexer.schema.monitoring import IndicesStats
@@ -23,7 +15,6 @@ from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_cl
15
from app.connectors.wazuh_indexer.utils.universal import format_indices_stats
16
from app.connectors.wazuh_indexer.utils.universal import format_node_allocation
17
from app.connectors.wazuh_indexer.utils.universal import format_shards
26
-from app.db.db_session import engine
18
19
20
def cluster_healthcheck() -> Union[ClusterHealthResponse, Dict[str, str]]:
backend/app/connectors/wazuh_indexer/utils/universal.py
-9
@@ -2,24 +2,15 @@ from datetime import datetime
2
from datetime import timedelta
3
from typing import Any
4
from typing import Dict
5
-from typing import Generator
5
from typing import Iterable
7
-from typing import List
6
from typing import Tuple
9
-from typing import Type
7
11
-import requests
8
from elasticsearch7 import Elasticsearch
9
from loguru import logger
14
-from sqlmodel import Session
15
-from sqlmodel import select
10
17
-from app.connectors.models import Connectors
18
-from app.connectors.schema import ConnectorResponse
11
from app.connectors.utils import get_connector_info_from_db
12
from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel
13
from app.connectors.wazuh_indexer.schema.indices import Indices
22
-from app.db.db_session import engine
14
15
16
def verify_wazuh_indexer_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
backend/app/connectors/wazuh_manager/models/rules.py
-3
@@ -1,10 +1,7 @@
1
import datetime
2
from typing import Optional
3
4
-from pydantic import EmailStr
5
-from pydantic import validator
4
from sqlmodel import Field
7
-from sqlmodel import Relationship
5
from sqlmodel import SQLModel
6
7
backend/app/connectors/wazuh_manager/schema/rules.py
-3
@@ -1,8 +1,5 @@
1
-from typing import Any
2
-from typing import Dict
1
from typing import List
2
from typing import Optional
5
-from typing import Union
3
4
from pydantic import BaseModel
5
backend/app/connectors/wazuh_manager/services/rules.py
-3
@@ -1,13 +1,10 @@
1
from typing import Any
2
from typing import Dict
3
from typing import List
4
-from typing import Optional
4
from typing import Tuple
5
from typing import Union
6
8
-import requests
7
import xmltodict
10
-from loguru import logger
8
9
from app.connectors.wazuh_manager.schema.rules import RuleDisable
10
from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse
backend/app/connectors/wazuh_manager/utils/universal.py
-5
@@ -4,13 +4,8 @@ from typing import Optional
4
5
import requests
6
from loguru import logger
7
-from sqlmodel import Session
8
-from sqlmodel import select
7
10
-from app.connectors.models import Connectors
11
-from app.connectors.schema import ConnectorResponse
8
from app.connectors.utils import get_connector_info_from_db
13
-from app.db.db_session import engine
9
10
11
def verify_wazuh_manager_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
backend/app/customers/routes/customers.py
-14
@@ -1,19 +1,10 @@
1
-from typing import Any
2
-from typing import Dict
3
-from typing import List
4
-from typing import Optional
5
-from typing import Type
6
-
1
from fastapi import APIRouter
8
-from fastapi import Depends
2
from fastapi import HTTPException
3
from fastapi import Query
11
-from fastapi import Security
4
from loguru import logger
5
from starlette.status import HTTP_401_UNAUTHORIZED
6
7
# App specific imports
16
-from app.auth.routes.auth import auth_handler
8
from app.customers.schema.customers import AgentModel
9
from app.customers.schema.customers import AgentsResponse
10
from app.customers.schema.customers import CustomerFullResponse
@@ -28,13 +19,8 @@ from app.db.universal_models import Customers
19
from app.db.universal_models import CustomersMeta
20
from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
21
from app.healthchecks.agents.schema.agents import AgentModel
31
-from app.healthchecks.agents.schema.agents import HostLogsSearchBody
32
-from app.healthchecks.agents.schema.agents import HostLogsSearchResponse
22
from app.healthchecks.agents.schema.agents import TimeCriteriaModel
34
-from app.healthchecks.agents.services.agents import host_logs
35
-from app.healthchecks.agents.services.agents import velociraptor_agent_healthcheck
23
from app.healthchecks.agents.services.agents import velociraptor_agents_healthcheck
37
-from app.healthchecks.agents.services.agents import wazuh_agent_healthcheck
24
from app.healthchecks.agents.services.agents import wazuh_agents_healthcheck
25
26
customers_router = APIRouter()
backend/app/db/all_models.py
+5
-5
@@ -1,6 +1,6 @@
1
# all_models.py
2
-from app.auth.models.users import User
3
-from app.connectors.models import Connectors
4
-from app.connectors.sublime.models.alerts import SublimeAlerts
5
-from app.connectors.wazuh_manager.models.rules import DisabledRule
6
-from app.db.universal_models import *
2
+# from app.auth.models.users import User
3
+# from app.connectors.models import Connectors
4
+# from app.connectors.sublime.models.alerts import SublimeAlerts
5
+# from app.connectors.wazuh_manager.models.rules import DisabledRule
6
+# from app.db.universal_models import Customers, CustomersMeta, Agents
backend/app/db/db_populate.py
-2
@@ -1,5 +1,3 @@
1
-from datetime import datetime
2
-
1
from loguru import logger
2
from sqlmodel import Session
3
backend/app/db/db_setup.py
+8
-7
@@ -3,14 +3,15 @@ from sqlalchemy import inspect
3
from sqlmodel import Session
4
from sqlmodel import SQLModel
5
6
-# from app.connectors.models import Connectors
7
-# from app.auth.models.users import User
8
-# from app.connectors.wazuh_manager.models.rules import DisabledRule
9
-# from app.agents.models.agents import Agents
10
-# from app.customers.models.customers import Customers
11
-# from app.connectors.sublime.models.alerts import SublimeAlerts
12
-from app.db.all_models import *
6
+# from app.db.all_models import *
7
+from app.auth.models.users import User
8
+from app.connectors.models import Connectors
9
+from app.connectors.sublime.models.alerts import SublimeAlerts
10
+from app.connectors.wazuh_manager.models.rules import DisabledRule
11
from app.db.db_populate import add_connectors_if_not_exist
12
+from app.db.universal_models import Agents
13
+from app.db.universal_models import Customers
14
+from app.db.universal_models import CustomersMeta
15
16
17
def create_tables(engine):
backend/app/db/universal_models.py
-1
@@ -1,7 +1,6 @@
1
from datetime import datetime
2
from typing import Optional
3
4
-from sqlmodel import Column
4
from sqlmodel import Field
5
from sqlmodel import Relationship
6
from sqlmodel import SQLModel
backend/app/healthchecks/agents/routes/agents.py
-9
@@ -1,21 +1,12 @@
1
-from datetime import datetime
2
-from typing import Dict
3
-from typing import List
4
-
1
from fastapi import APIRouter
6
-from fastapi import Depends
2
from fastapi import HTTPException
3
from fastapi import Query
9
-from fastapi import Security
4
from loguru import logger
5
from starlette.status import HTTP_401_UNAUTHORIZED
6
13
-# App specific imports
14
-from app.auth.routes.auth import auth_handler
7
from app.db.db_session import session
8
from app.db.universal_models import Agents
9
from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
18
-from app.healthchecks.agents.schema.agents import AgentModel
10
from app.healthchecks.agents.schema.agents import HostLogsSearchBody
11
from app.healthchecks.agents.schema.agents import HostLogsSearchResponse
12
from app.healthchecks.agents.schema.agents import TimeCriteriaModel
backend/app/healthchecks/agents/schema/agents.py
-1
@@ -1,5 +1,4 @@
1
from datetime import datetime
2
-from enum import Enum
2
from typing import Any
3
from typing import Dict
4
from typing import List
backend/app/healthchecks/agents/services/agents.py
-15
@@ -1,28 +1,14 @@
1
from datetime import datetime
2
from datetime import timedelta
3
-from typing import Any
4
-from typing import Dict
5
-from typing import List
3
from typing import Optional
4
from typing import Type
5
9
-from fastapi import APIRouter
10
-from fastapi import Depends
6
from fastapi import HTTPException
12
-from fastapi import Security
7
from loguru import logger
14
-from starlette.status import HTTP_401_UNAUTHORIZED
8
16
-# App specific imports
17
-from app.auth.routes.auth import auth_handler
9
from app.connectors.wazuh_indexer.utils.universal import LogsQueryBuilder
10
from app.connectors.wazuh_indexer.utils.universal import collect_indices
11
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
21
-from app.connectors.wazuh_indexer.utils.universal import format_indices_stats
22
-from app.connectors.wazuh_indexer.utils.universal import format_node_allocation
23
-from app.connectors.wazuh_indexer.utils.universal import format_shards
24
-from app.db.db_session import session
25
-from app.db.universal_models import Agents
12
from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
13
from app.healthchecks.agents.schema.agents import AgentModel
14
from app.healthchecks.agents.schema.agents import CollectLogsResponse
@@ -30,7 +16,6 @@ from app.healthchecks.agents.schema.agents import ExtendedAgentModel
16
from app.healthchecks.agents.schema.agents import HostLogsSearchBody
17
from app.healthchecks.agents.schema.agents import HostLogsSearchResponse
18
from app.healthchecks.agents.schema.agents import LogsSearchBody
33
-from app.healthchecks.agents.schema.agents import LogsSearchResponse
19
from app.healthchecks.agents.schema.agents import TimeCriteriaModel
20
21
backend/app/integrations/alert_escalation/routes/general_alert.py
-10
@@ -1,16 +1,6 @@
1
-from typing import List
2
-
1
from fastapi import APIRouter
4
-from fastapi import Depends
5
-from fastapi import HTTPException
6
-from fastapi import Security
2
from loguru import logger
8
-from starlette.status import HTTP_401_UNAUTHORIZED
3
10
-# App specific imports
11
-from app.auth.routes.auth import auth_handler
12
-from app.connectors.wazuh_indexer.utils.universal import collect_indices
13
-from app.db.db_session import session
4
from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
5
from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse
6
from app.integrations.alert_escalation.services.general_alert import create_alert
backend/app/integrations/alert_escalation/schema/general_alert.py
-2
@@ -3,12 +3,10 @@ from typing import Any
3
from typing import Dict
4
from typing import List
5
from typing import Optional
6
-from typing import Union
6
7
from pydantic import BaseModel
8
from pydantic import Extra
9
from pydantic import Field
11
-from pydantic import validator
10
11
12
class ValidIocFields(Enum):
backend/app/integrations/alert_escalation/services/general_alert.py
+1
-23
@@ -1,34 +1,12 @@
1
-from typing import Any
2
-from typing import Dict
3
-from typing import List
1
from typing import Optional
2
from typing import Set
6
-from typing import Type
7
-from typing import Union
3
9
-from dfir_iris_client.alert import Alert
10
-from elasticsearch7 import Elasticsearch
4
from fastapi import HTTPException
5
from loguru import logger
13
-from pydantic import BaseModel
14
-from sqlmodel import Session
15
-from sqlmodel import select
6
17
-from app.connectors.dfir_iris.utils.universal import create_dfir_iris_client
18
-from app.connectors.dfir_iris.utils.universal import fetch_and_parse_data
7
from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
8
from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
21
-from app.connectors.models import Connectors
22
-from app.connectors.schema import ConnectorResponse
23
-from app.connectors.utils import get_connector_info_from_db
24
-from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel
25
-from app.connectors.wazuh_indexer.utils.universal import AlertsQueryBuilder
26
-from app.connectors.wazuh_indexer.utils.universal import collect_indices
9
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
28
-from app.connectors.wazuh_indexer.utils.universal import format_indices_stats
29
-from app.connectors.wazuh_indexer.utils.universal import format_node_allocation
30
-from app.connectors.wazuh_indexer.utils.universal import format_shards
31
-from app.db.db_session import engine
10
from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
11
from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse
12
from app.integrations.alert_escalation.schema.general_alert import GenericAlertModel
@@ -118,7 +96,7 @@ def build_alert_payload(alert_details: GenericAlertModel, agent_data, ioc_payloa
96
alert_iocs=[ioc_payload],
97
)
98
else:
121
- logger.info(f"Alert does not have IoC")
99
+ logger.info("Alert does not have IoC")
100
return IrisAlertPayload(
101
alert_title=alert_details._source.rule_description,
102
alert_description=alert_details._source.rule_description,
backend/app/integrations/alert_escalation/utils/universal.py
-16
@@ -1,35 +1,19 @@
1
import ipaddress
2
import re
3
from abc import ABC
4
-from datetime import datetime
5
-from datetime import timedelta
4
from typing import Any
5
from typing import Dict
8
-from typing import Generator
9
-from typing import Iterable
10
-from typing import List
6
from typing import Optional
12
-from typing import Tuple
13
-from typing import Type
7
from typing import Union
8
9
import regex
17
-import requests
10
from elasticsearch7 import Elasticsearch
11
from fastapi import HTTPException
12
from loguru import logger
21
-from sqlmodel import Session
22
-from sqlmodel import select
13
24
-from app.connectors.models import Connectors
25
-from app.connectors.schema import ConnectorResponse
14
from app.connectors.utils import get_connector_info_from_db
27
-from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel
28
-from app.connectors.wazuh_indexer.schema.indices import Indices
15
from app.db.all_models import Agents
30
-from app.db.db_session import engine
16
from app.db.db_session import session
32
-from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
17
from app.healthchecks.agents.schema.agents import AgentModel
18
19
backend/app/integrations/dnstwist/routes/analyze.py
-6
@@ -2,18 +2,12 @@ import regex
2
from fastapi import APIRouter
3
from fastapi import Depends
4
from fastapi import HTTPException
5
-from fastapi import Security
6
-from fastapi import security
7
-from fastapi.security import HTTPAuthorizationCredentials
5
from loguru import logger
6
10
-from app.auth.utils import AuthHandler
7
from app.db.db_session import session
8
from app.integrations.dnstwist.schema.analyze import DomainAnalysisResponse
9
from app.integrations.dnstwist.schema.analyze import DomainRequestBody
10
from app.integrations.dnstwist.services.analyze import analyze_domain
15
-from app.integrations.dnstwist.services.analyze import analyze_domain_phishing
16
-from app.smtp.schema.configure import SMTPResponse
11
12
dnstwist_router = APIRouter()
13
backend/app/integrations/dnstwist/services/analyze.py
-3
@@ -1,6 +1,3 @@
1
-from typing import Any
2
-from typing import Optional
3
-
1
import dnstwist
2
from loguru import logger
3
backend/app/smtp/routes/configure.py
-6
@@ -1,15 +1,9 @@
1
from fastapi import APIRouter
2
-from fastapi import Depends
2
from fastapi import HTTPException
4
-from fastapi import Security
5
-from fastapi import security
6
-from fastapi.security import HTTPAuthorizationCredentials
3
from loguru import logger
4
5
from app.auth.models.users import SMTP
6
from app.auth.models.users import SMTPInput
11
-from app.auth.models.users import User
12
-from app.auth.services.universal import find_user
7
from app.auth.services.universal import select_all_users
8
from app.auth.utils import AuthHandler
9
from app.db.db_session import session
backend/app/smtp/routes/reports.py
+1
-6
@@ -1,21 +1,16 @@
1
from fastapi import APIRouter
2
-from fastapi import Depends
2
from fastapi import HTTPException
4
-from fastapi import Security
5
-from fastapi import security
6
-from fastapi.security import HTTPAuthorizationCredentials
3
from loguru import logger
4
5
from app.auth.models.users import SMTP
6
from app.auth.models.users import SMTPInput
11
-from app.auth.models.users import User
12
-from app.auth.services.universal import find_user
7
from app.auth.services.universal import select_all_users
8
from app.auth.utils import AuthHandler
9
from app.db.db_session import session
10
from app.smtp.schema.configure import SMTPResponse
11
12
smtp_reports_router = APIRouter()
13
+auth_handler = AuthHandler()
14
15
16
# ! TODO: Add SMTP reporting all things. Example is in the services/reports.py and services/create_report.py file
backend/copilot.py
-3
@@ -4,8 +4,6 @@ from fastapi import HTTPException
4
from fastapi import Request
5
from fastapi.middleware.cors import CORSMiddleware
6
from fastapi.responses import JSONResponse
7
-from sqlmodel import SQLModel
8
-from sqlmodel import create_engine
7
8
from app.agents.routes.agents import agents_router
9
from app.auth.routes.auth import user_router
@@ -39,7 +37,6 @@ from app.integrations.alert_escalation.routes.general_alert import (
37
)
38
from app.integrations.dnstwist.routes.analyze import dnstwist_router
39
from app.smtp.routes.configure import smtp_router
42
-from settings import SQLALCHEMY_DATABASE_URI
40
41
app = FastAPI(description="CoPilot API", version="0.1.0", title="CoPilot API")
42