
from typing import Dict, Optional, Any
from pathlib import Path
from datetime import datetime, timezone
import json
import requests
import hmac
import hashlib
import base64
from pentestmcp import config
"""bloodhound api calls automation in prog....."""
config.load_variables()
BHE_SCHEME = "http"
PRINT_PRINCIPALS = False
PRINT_ATTACK_PATH_TIMELINE_DATA = False
PRINT_POSTURE_DATA = False

DATA_START = "1970-01-01T00:00:00.000Z"
DATA_END = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3] + 'Z' # Now

class Credentials(object):
    def __init__(self, token_id: str, token_key: str) -> None:
        self.token_id = token_id
        self.token_key = token_key

class Domain(object):
    def __init__(self, name: str, id: str, collected: bool, domain_type: str) -> None:
        self.name = name
        self.id = id
        self.type = domain_type
        self.collected = collected

class APIVersion(object):
    def __init__(self, api_version: str, server_version: str) -> None:
        self.api_version = api_version
        self.server_version = server_version
class Client(object):
    def __init__(self, scheme: str, host: str, port: int, credentials: Credentials) -> None:
        self._scheme = scheme
        self._host = host
        self._port = port
        self._credentials = credentials

    def _format_url(self, uri: str) -> str:
        formatted_uri = uri
        if uri.startswith("/"):
            formatted_uri = formatted_uri[1:]

        return f"{self._scheme}://{self._host}:{self._port}/{formatted_uri}"

    def _request(
        self,
        method: str,
        uri: str,
        body: Optional[bytes] = None,
        content_type: str =None ,
        timeout: Optional[float] = None,
    ) -> requests.Response:
        """
        Signed request used by BloodHound CE.
        - `content_type` lets us upload application/zip or application/json.
        - Signature header is base64 text (not bytes).
        """
        # start chain with token_key as HMAC key
        digester = hmac.new(self._credentials.token_key.encode("utf-8"), None, hashlib.sha256)

        # OperationKey: method + uri (no delimiter)
        digester.update(f"{method}{uri}".encode("utf-8"))

        # chain
        digester = hmac.new(digester.digest(), None, hashlib.sha256)

        # RequestDate (RFC3339-like). Use UTC and include 'Z' to match expectations.
        datetime_formatted = datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")
        # docs say chain uses hours only (first 13 chars: YYYY-MM-DDTHH)
        digester.update(datetime_formatted[:13].encode("utf-8"))

        # chain again
        digester = hmac.new(digester.digest(), None, hashlib.sha256)

        # include body bytes in signature if present
        if body is not None:
            digester.update(body)

        # signature as BASE64 string (important: decode to ascii)
        signature_b64 = base64.b64encode(digester.digest()).decode("ascii")
        if content_type != None:
            
            headers: Dict[str, str] = {
                "User-Agent": "bhe-python-sdk 0001",
                "Authorization": f"bhesignature {self._credentials.token_id}",
                "RequestDate": datetime_formatted,
                "Signature": signature_b64,
                "Content-Type": content_type,
            }
        else:
            headers: Dict[str, str] = {
            "User-Agent": "bhe-python-sdk 0001",
            "Authorization": f"bhesignature {self._credentials.token_id}",
            "RequestDate": datetime_formatted,
            "Signature": signature_b64,
        }

        # perform request (use data=body so we pass raw bytes as-is)
        resp = requests.request(method=method, url=self._format_url(uri), headers=headers, data=body, timeout=timeout)
        resp.raise_for_status()  # raise on HTTP error (optional; useful for automation)
        return resp

    def get_version(self) -> APIVersion:
        response = self._request("GET", "/api/version")
        payload = response.json()

        return APIVersion(api_version=payload["data"]["API"]["current_version"], server_version=payload["data"]["server_version"])


def connectToApi():
    # This might be best loaded from a file
    credentials = Credentials(
        token_id=config.BHE_TOKEN_ID,
        token_key=config.BHE_TOKEN_KEY,
    )

    client = Client(scheme=BHE_SCHEME, host=config.BHE_DOMAIN, port=config.BHE_PORT, credentials=credentials)
    
    return client.get_version().api_version



def upload_zip(zip_path: str) -> Dict[str, Any]:
    """
    Upload a SharpHound zip (or BloodHound JSON zip) using the BHCE file-upload job flow:
      1) POST /api/v2/file-upload/start  -> get upload id
      2) POST /api/v2/file-upload/{id}   -> upload bytes (Content-Type: application/zip)
      3) POST /api/v2/file-upload/{id}/end -> finish job

    Returns the final job response JSON (or raises on error).
    """


    # create client
    credentials = Credentials(token_id=config.BHE_TOKEN_ID, token_key=config.BHE_TOKEN_KEY)
    client = Client(scheme=BHE_SCHEME, host=config.BHE_DOMAIN, port=config.BHE_PORT, credentials=credentials)

    p = Path(zip_path)
    if not p.exists():
        raise FileNotFoundError(f"{zip_path} not found")

    zip_bytes = p.read_bytes()

    # 1) create upload job
    resp = client._request("POST", "/api/v2/file-upload/start", content_type="application/json")
    # response body should contain the upload job id in JSON: { "data": { "id": "<id>" }, ... }
    start_payload = resp.json()
    upload_id = start_payload.get("data", {}).get("id")
    if not upload_id:
        raise RuntimeError(f"Could not create upload job: {start_payload}")

    # 2) upload the zip bytes (include body bytes when signing)
    upload_uri = f"/api/v2/file-upload/{upload_id}"
    client._request("POST", upload_uri, body=zip_bytes, content_type="application/zip")
    # 3) finish the job
    end_uri = f"/api/v2/file-upload/{upload_id}/end"
    end_resp = client._request("POST", end_uri, body=None, content_type="application/json")
    # API returns 204 No Content on success; depending on your client .json() may fail.
    if end_resp.status_code == 204:
        return {"status": "accepted", "upload_id": upload_id}
    else:
        return end_resp

def list_saved_queries():
    credentials = Credentials(
        token_id=config.BHE_TOKEN_ID,
        token_key=config.BHE_TOKEN_KEY,
    )

    client = Client(scheme=BHE_SCHEME, host=config.BHE_DOMAIN, port=config.BHE_PORT, credentials=credentials)
    response=client._request("GET","/api/v2/saved-queries")
    return response.json()
def run_cypher_query(query):
    credentials = Credentials(
        token_id=config.BHE_TOKEN_ID,
        token_key=config.BHE_TOKEN_KEY,
    )

    client = Client(scheme=BHE_SCHEME, host=config.BHE_DOMAIN, port=config.BHE_PORT, credentials=credentials)
    payload={"query":query}
    data=json.dumps(payload).encode('utf-8')
    response=client._request("POST","/api/v2/graphs/cypher",body=data,content_type="application/json")
    return response.json()

"""NEO4J implementation coming next """
