# Contract pinned from MythicMeta/Mythic_Scripting master @ 2026-06-10 (raw.githubusercontent.com/MythicMeta/Mythic_Scripting/master/mythic/mythic.py)
"""Mythic 3.x C2 adapter.

Transport: POST https://<host>:7443/graphql
Header: apitoken: <token>
Backend: Hasura-proxied Postgres behind nginx.

M1: test_connection()
M2: list_callbacks(), create_task()
M3: get_task(), get_task_output()
M4: list_callback_tasks()
"""
from __future__ import annotations

import requests

from backend.app.services.c2.adapter import (
    C2Adapter,
    C2Callback,
    C2Error,
    C2Health,
    C2TaskPage,
    C2TaskStatus,
)

_HEALTH_QUERY = "{ __typename }"

_CALLBACKS_QUERY = """
query {
  callback(order_by: {id: asc}, where: {active: {_eq: true}}) {
    id
    display_id
    active
    host
    user
    domain
    last_checkin
  }
}
"""

_CREATE_TASK_MUTATION = """
mutation CreateTask($callback_id: Int!, $command: String!, $params: String!) {
  createTask(
    callback_id: $callback_id,
    command: $command,
    params: $params,
    tasking_location: "command_line"
  ) {
    id
    display_id
    error
  }
}
"""


class MythicAdapter(C2Adapter):
    """Real Mythic 3.x adapter using GraphQL over HTTP."""

    def __init__(self, url: str, api_token: str, verify_tls: bool = True) -> None:
        self._url = url.rstrip("/") + "/graphql"
        self._token = api_token
        self._verify = verify_tls

    def _headers(self) -> dict[str, str]:
        return {
            "Content-Type": "application/json",
            "apitoken": self._token,
        }

    def _post(self, body: dict) -> dict:
        resp = requests.post(
            self._url,
            json=body,
            headers=self._headers(),
            verify=self._verify,
            timeout=10,
            allow_redirects=False,
        )
        resp.raise_for_status()
        return resp.json()

    def test_connection(self) -> C2Health:
        """POST a trivial introspection query to verify reachability and token validity."""
        try:
            resp = requests.post(
                self._url,
                json={"query": _HEALTH_QUERY},
                headers=self._headers(),
                verify=self._verify,
                timeout=10,
                allow_redirects=False,
            )
            if resp.status_code == 200:
                return C2Health(ok=True)
            return C2Health(ok=False, error=f"HTTP {resp.status_code}")
        except requests.RequestException as exc:
            return C2Health(ok=False, error=str(exc))

    def list_callbacks(self) -> list[C2Callback]:
        """Return active callbacks from Mythic (filtered server-side: active=true)."""
        try:
            data = self._post({"query": _CALLBACKS_QUERY})
        except requests.RequestException as exc:
            raise C2Error(str(exc)) from exc

        callbacks_raw = data.get("data", {}).get("callback", [])
        return [
            C2Callback(
                display_id=cb["display_id"],
                active=cb["active"],
                host=cb.get("host") or "",
                user=cb.get("user") or "",
                domain=cb.get("domain") or "",
                last_checkin=cb.get("last_checkin") or "",
            )
            for cb in callbacks_raw
        ]

    def create_task(
        self,
        callback_display_id: int,
        command: str,
        params: str | None = None,
    ) -> int:
        """Issue a task on a callback; return Mythic task display_id."""
        try:
            data = self._post({
                "query": _CREATE_TASK_MUTATION,
                "variables": {
                    "callback_id": callback_display_id,
                    "command": command,
                    "params": params or "",
                },
            })
        except requests.RequestException as exc:
            raise C2Error(str(exc)) from exc

        task_data = data.get("data", {}).get("createTask", {})
        error_msg = task_data.get("error")
        if error_msg:
            raise C2Error(error_msg)
        return int(task_data["display_id"])

    def get_task(self, task_display_id: int) -> C2TaskStatus:
        raise NotImplementedError("M3")

    def get_task_output(self, task_display_id: int) -> str:
        raise NotImplementedError("M3")

    def list_callback_tasks(
        self,
        callback_display_id: int,
        page: int = 1,
        page_size: int = 25,
    ) -> C2TaskPage:
        raise NotImplementedError("M4")