precommit -fixes
taylorwalton committed
Feb 19, 2025 at 20:48 UTC
85e84ca9e3921522a201a7ee2d112d1050adadb3
8 files changed
+51
-54
backend/app/connectors/portainer/routes/portainer.py
+3
-2
@@ -3,12 +3,13 @@ from fastapi import Security
3
from loguru import logger
4
5
from app.auth.utils import AuthHandler
6
-#from app.connectors.portainer.schema.integrations import ExecuteWorkflowRequest
6
+
7
+# from app.connectors.portainer.schema.integrations import ExecuteWorkflowRequest
8
from app.connectors.portainer.schema.nodes import NodesResponse
9
+from app.connectors.portainer.schema.stack import StackResponse
10
from app.connectors.portainer.services.nodes import get_node_details
11
from app.connectors.portainer.services.stack import create_wazuh_customer_stack
12
from app.customer_provisioning.schema.provision import ProvisionNewCustomer
11
-from app.connectors.portainer.schema.stack import StackResponse
13
14
portainer_integrations_router = APIRouter()
15
backend/app/connectors/portainer/schema/nodes.py
+16
-3
@@ -1,52 +1,65 @@
1
-from typing import List, Optional
2
-from pydantic import BaseModel
1
from datetime import datetime
2
+from typing import List
3
+from typing import Optional
4
+
5
+from pydantic import BaseModel
6
+
7
8
class Version(BaseModel):
9
Index: int
10
11
+
12
class NodeSpec(BaseModel):
13
Labels: dict
14
Role: str
15
Availability: str
16
17
+
18
class Platform(BaseModel):
19
Architecture: str
20
OS: str
21
22
+
23
class Resources(BaseModel):
24
NanoCPUs: int
25
MemoryBytes: int
26
27
+
28
class Plugin(BaseModel):
29
Type: str
30
Name: str
31
32
+
33
class TLSInfo(BaseModel):
34
TrustRoot: Optional[str] = None
35
CertIssuerSubject: Optional[str] = None
36
CertIssuerPublicKey: Optional[str] = None
37
38
+
39
class Engine(BaseModel):
40
EngineVersion: str
41
Plugins: List[Plugin]
42
TLSInfo: Optional[TLSInfo] = None
43
44
+
45
class Description(BaseModel):
46
Hostname: str
47
Platform: Platform
48
Resources: Resources
49
Engine: Engine
50
51
+
52
class ManagerStatusResults(BaseModel):
53
Leader: bool
54
Reachability: str
55
Addr: str
56
57
+
58
class Status(BaseModel):
59
State: str
60
Addr: str
61
62
+
63
class Node(BaseModel):
64
ID: str
65
Version: Version
@@ -57,8 +70,8 @@ class Node(BaseModel):
70
Status: Status
71
ManagerStatus: Optional[ManagerStatusResults] = None
72
73
+
74
class NodesResponse(BaseModel):
75
nodes: List[Node]
76
success: bool
77
message: str
64
-
backend/app/connectors/portainer/schema/stack.py
+4
-2
@@ -1,6 +1,8 @@
1
-from typing import List, Optional, Dict, Any
1
+from typing import Any
2
+from typing import List
3
+from typing import Optional
4
+
5
from pydantic import BaseModel
3
-from datetime import datetime
6
7
8
class ResourceControl(BaseModel):
backend/app/connectors/portainer/services/nodes.py
+2
-2
@@ -1,8 +1,8 @@
1
from loguru import logger
2
3
from app.connectors.portainer.schema.nodes import NodesResponse
4
-
5
-from app.connectors.portainer.utils.universal import send_get_request, get_endpoint_id
4
+from app.connectors.portainer.utils.universal import get_endpoint_id
5
+from app.connectors.portainer.utils.universal import send_get_request
6
7
8
async def get_node_details() -> NodesResponse:
backend/app/connectors/portainer/services/stack.py
+10
-8
@@ -1,9 +1,13 @@
1
-from loguru import logger
1
from pathlib import Path
3
-from app.connectors.portainer.utils.universal import send_post_request, get_endpoint_id, get_swarm_id
2
+
3
+from loguru import logger
4
+
5
from app.agents.routes.agents import get_wazuh_manager_version
5
-from app.customer_provisioning.schema.provision import ProvisionNewCustomer
6
from app.connectors.portainer.schema.stack import StackResponse
7
+from app.connectors.portainer.utils.universal import get_endpoint_id
8
+from app.connectors.portainer.utils.universal import get_swarm_id
9
+from app.connectors.portainer.utils.universal import send_post_request
10
+from app.customer_provisioning.schema.provision import ProvisionNewCustomer
11
12
13
async def create_wazuh_customer_stack(request: ProvisionNewCustomer) -> StackResponse:
@@ -22,9 +26,9 @@ async def create_wazuh_customer_stack(request: ProvisionNewCustomer) -> StackRes
26
# Replace the placeholders in the template with the actual values
27
template = template.replace("{{ wazuh_worker_customer_code }}", formatted_customer_name)
28
template = template.replace("{{ wazuh_manager_version }}", wazuh_manager_version)
25
- template = template.replace('REPLACE_LOG', request.wazuh_logs_port)
26
- template = template.replace('REPLACE_REGISTRATION', request.wazuh_registration_port)
27
- template = template.replace('REPLACE_API', request.wazuh_api_port)
29
+ template = template.replace("REPLACE_LOG", request.wazuh_logs_port)
30
+ template = template.replace("REPLACE_REGISTRATION", request.wazuh_registration_port)
31
+ template = template.replace("REPLACE_API", request.wazuh_api_port)
32
logger.info(f"Template: {template}")
33
34
endpoint_id = await get_endpoint_id()
@@ -42,5 +46,3 @@ async def create_wazuh_customer_stack(request: ProvisionNewCustomer) -> StackRes
46
response = await send_post_request(endpoint=create_stack_url, data=payload)
47
logger.info(f"Response: {response}")
48
return StackResponse(**response)
45
-
46
-
backend/app/connectors/portainer/utils/universal.py
+13
-35
@@ -1,11 +1,12 @@
1
from typing import Any
2
from typing import Dict
3
from typing import Optional
4
+from urllib.parse import urljoin
5
6
import requests
7
from fastapi import HTTPException
8
from loguru import logger
8
-from urllib.parse import urljoin
9
+
10
from app.connectors.utils import get_connector_info_from_db
11
from app.db.db_session import get_db_session
12
@@ -24,6 +25,7 @@ async def get_endpoint_id() -> int:
25
return int(endpoint["Id"])
26
return None
27
28
+
29
async def get_swarm_id() -> int:
30
"""
31
Returns the ID of the swarm.
@@ -40,21 +42,14 @@ async def get_portainer_jwt() -> str:
42
"""Get JWT token from Portainer API."""
43
logger.info("Getting portainer authentication token")
44
async with get_db_session() as session: # This will correctly enter the context manager
43
- attributes = await get_connector_info_from_db('Portainer', session)
45
+ attributes = await get_connector_info_from_db("Portainer", session)
46
logger.info(f"Attributes: {attributes}")
47
try:
48
auth_endpoint = urljoin(attributes["connector_url"], "/api/auth")
49
48
- auth_payload = {
49
- "username": attributes["connector_username"],
50
- "password": attributes["connector_password"]
51
- }
50
+ auth_payload = {"username": attributes["connector_username"], "password": attributes["connector_password"]}
51
53
- response = requests.post(
54
- auth_endpoint,
55
- json=auth_payload,
56
- verify=False # If using self-signed cert
57
- )
52
+ response = requests.post(auth_endpoint, json=auth_payload, verify=False) # If using self-signed cert
53
54
response.raise_for_status()
55
# The JWT token is in response.json()["jwt"]
@@ -92,16 +87,9 @@ async def verify_portainer_credentials(attributes: Dict[str, Any]) -> Dict[str,
87
# If API key auth fails, try JWT authentication
88
if portainer_apps.status_code != 200:
89
auth_endpoint = urljoin(attributes["connector_url"], "/api/auth")
95
- auth_payload = {
96
- "username": attributes["connector_username"],
97
- "password": attributes["connector_password"]
98
- }
90
+ auth_payload = {"username": attributes["connector_username"], "password": attributes["connector_password"]}
91
100
- jwt_response = requests.post(
101
- auth_endpoint,
102
- json=auth_payload,
103
- verify=False
104
- )
92
+ jwt_response = requests.post(auth_endpoint, json=auth_payload, verify=False)
93
94
if jwt_response.status_code == 200:
95
jwt_token = jwt_response.json()["jwt"]
@@ -110,26 +98,16 @@ async def verify_portainer_credentials(attributes: Dict[str, Any]) -> Dict[str,
98
"connectionSuccessful": True,
99
"message": "Portainer connection successful via JWT",
100
"authMethod": "jwt",
113
- "jwt": jwt_token
101
+ "jwt": jwt_token,
102
}
103
else:
116
- logger.error(
117
- f"Both API key and JWT authentication failed. JWT error: {jwt_response.text}"
118
- )
119
- return {
120
- "connectionSuccessful": False,
121
- "message": "Both API key and JWT authentication failed",
122
- "authMethod": None
123
- }
104
+ logger.error(f"Both API key and JWT authentication failed. JWT error: {jwt_response.text}")
105
+ return {"connectionSuccessful": False, "message": "Both API key and JWT authentication failed", "authMethod": None}
106
107
logger.info(
108
f"Connection to {attributes['connector_url']} successful via API key",
109
)
128
- return {
129
- "connectionSuccessful": True,
130
- "message": "Portainer connection successful via API key",
131
- "authMethod": "api_key"
132
- }
110
+ return {"connectionSuccessful": True, "message": "Portainer connection successful via API key", "authMethod": "api_key"}
111
112
except Exception as e:
113
logger.error(
@@ -138,7 +116,7 @@ async def verify_portainer_credentials(attributes: Dict[str, Any]) -> Dict[str,
116
return {
117
"connectionSuccessful": False,
118
"message": f"Connection to {attributes['connector_url']} failed with error: {e}",
141
- "authMethod": None
119
+ "authMethod": None,
120
}
121
122
backend/app/connectors/services.py
+2
-1
@@ -17,8 +17,8 @@ from app.connectors.cortex.utils.universal import verify_cortex_connection
17
from app.connectors.grafana.utils.universal import verify_grafana_connection
18
from app.connectors.graylog.utils.universal import verify_graylog_connection
19
from app.connectors.influxdb.utils.universal import verify_influxdb_connection
20
-from app.connectors.portainer.utils.universal import verify_portainer_connection
20
from app.connectors.models import Connectors
21
+from app.connectors.portainer.utils.universal import verify_portainer_connection
22
from app.connectors.schema import ConnectorResponse
23
from app.connectors.shuffle.utils.universal import verify_shuffle_connection
24
from app.connectors.sublime.utils.universal import verify_sublime_connection
@@ -179,6 +179,7 @@ class VirustotalService(ConnectorServiceInterface):
179
) -> Optional[ConnectorResponse]:
180
return await verify_virustotal_connection(connector.connector_name)
181
182
+
183
# Portainer Service
184
class PortainerService(ConnectorServiceInterface):
185
async def verify_authentication(
backend/copilot.py
+1
-1
@@ -61,8 +61,8 @@ from app.routers import monitoring_alert
61
from app.routers import network_connectors
62
from app.routers import nuclei
63
from app.routers import office365
64
-from app.routers import sap_siem
64
from app.routers import portainer
65
+from app.routers import sap_siem
66
from app.routers import scheduler
67
from app.routers import scoutsuite
68
from app.routers import shuffle