Logan
255 lines
9.7 KiB
Python
255 lines
9.7 KiB
Python
import httpx
|
|
from typing import Optional
|
|
from app.config import settings
|
|
from app.internal.logger import logger
|
|
|
|
|
|
class C2Client:
|
|
def __init__(self):
|
|
self.base = settings.c2_url.rstrip("/")
|
|
|
|
def _headers(self) -> dict:
|
|
if settings.c2_service_key:
|
|
return {"Authorization": f"Bearer {settings.c2_service_key}"}
|
|
return {}
|
|
|
|
async def get_nodes(self) -> list:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.get(f"{self.base}/nodes", headers=self._headers())
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 get_nodes failed: {e}")
|
|
return []
|
|
|
|
async def get_systems(self) -> list:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.get(f"{self.base}/systems", headers=self._headers())
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 get_systems failed: {e}")
|
|
return []
|
|
|
|
async def get_tokens(self) -> list:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.get(f"{self.base}/tokens", headers=self._headers())
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 get_tokens failed: {e}")
|
|
return []
|
|
|
|
async def send_command(self, node_id: str, payload: dict) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/nodes/{node_id}/command",
|
|
json=payload,
|
|
headers=self._headers(),
|
|
)
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 send_command failed: {e}")
|
|
return False
|
|
|
|
async def find_node_for_system(self, system_id: str) -> Optional[dict]:
|
|
"""Return the first online node assigned to the given system."""
|
|
nodes = await self.get_nodes()
|
|
for node in nodes:
|
|
if (
|
|
node.get("assigned_system_id") == system_id
|
|
and node.get("status") in ("online", "recording")
|
|
):
|
|
return node
|
|
return None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Trips
|
|
# ------------------------------------------------------------------
|
|
|
|
async def get_trips(self) -> list:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.get(f"{self.base}/trips", headers=self._headers())
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 get_trips failed: {e}")
|
|
return []
|
|
|
|
async def get_trip(self, trip_id: str) -> Optional[dict]:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.get(f"{self.base}/trips/{trip_id}", headers=self._headers())
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 get_trip failed: {e}")
|
|
return None
|
|
|
|
async def create_trip(self, payload: dict) -> Optional[dict]:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(f"{self.base}/trips", json=payload, headers=self._headers())
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 create_trip failed: {e}")
|
|
return None
|
|
|
|
async def delete_trip(self, trip_id: str) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.delete(f"{self.base}/trips/{trip_id}", headers=self._headers())
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 delete_trip failed: {e}")
|
|
return False
|
|
|
|
async def invite_to_trip(self, trip_id: str, discord_user_id: str) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/trips/{trip_id}/invite/{discord_user_id}",
|
|
headers=self._headers(),
|
|
)
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 invite_to_trip failed: {e}")
|
|
return False
|
|
|
|
async def set_trip_visibility(self, trip_id: str, visibility: str) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.put(
|
|
f"{self.base}/trips/{trip_id}/visibility",
|
|
json={"visibility": visibility},
|
|
headers=self._headers(),
|
|
)
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 set_trip_visibility failed: {e}")
|
|
return False
|
|
|
|
async def link_discord_account(self, code: str, discord_user_id: str, discord_username: str) -> dict:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/auth/link",
|
|
json={"code": code, "discord_user_id": discord_user_id, "discord_username": discord_username},
|
|
headers=self._headers(),
|
|
)
|
|
if r.status_code == 404:
|
|
return {"error": "invalid_code"}
|
|
if r.status_code == 410:
|
|
return {"error": "expired"}
|
|
if r.status_code == 409:
|
|
return {"error": "already_linked"}
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 link_discord_account failed: {e}")
|
|
return {"error": "failed"}
|
|
|
|
async def join_trip(self, trip_id: str, user_id: str, username: str) -> bool | str:
|
|
"""Returns True on success, 'private' on 403, False on other errors."""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/trips/{trip_id}/join",
|
|
json={"discord_user_id": user_id, "discord_username": username},
|
|
headers=self._headers(),
|
|
)
|
|
if r.status_code == 403:
|
|
return "private"
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 join_trip failed: {e}")
|
|
return False
|
|
|
|
async def leave_trip(self, trip_id: str, user_id: str) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/trips/{trip_id}/leave",
|
|
json={"discord_user_id": user_id},
|
|
headers=self._headers(),
|
|
)
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 leave_trip failed: {e}")
|
|
return False
|
|
|
|
async def create_trip_event(self, trip_id: str, payload: dict) -> Optional[dict]:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/trips/{trip_id}/events",
|
|
json=payload,
|
|
headers=self._headers(),
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
logger.error(f"C2 create_trip_event failed: {e}")
|
|
return None
|
|
|
|
async def delete_trip_event(self, trip_id: str, event_id: str) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.delete(
|
|
f"{self.base}/trips/{trip_id}/events/{event_id}",
|
|
headers=self._headers(),
|
|
)
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 delete_trip_event failed: {e}")
|
|
return False
|
|
|
|
async def join_trip_event(
|
|
self, trip_id: str, event_id: str, user_id: str, username: str
|
|
) -> bool | str:
|
|
"""Returns True on success, 'not_on_trip' on 403, False on other errors."""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/trips/{trip_id}/events/{event_id}/join",
|
|
json={"discord_user_id": user_id, "discord_username": username},
|
|
headers=self._headers(),
|
|
)
|
|
if r.status_code == 403:
|
|
return "not_on_trip"
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 join_trip_event failed: {e}")
|
|
return False
|
|
|
|
async def leave_trip_event(self, trip_id: str, event_id: str, user_id: str) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
r = await client.post(
|
|
f"{self.base}/trips/{trip_id}/events/{event_id}/leave",
|
|
json={"discord_user_id": user_id},
|
|
headers=self._headers(),
|
|
)
|
|
r.raise_for_status()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"C2 leave_trip_event failed: {e}")
|
|
return False
|
|
|
|
|
|
c2 = C2Client()
|