envied/unshackle/vaults/HTTPAPI.py

import json
from typing import Iterator, Optional, Union
from uuid import UUID
from enum import Enum

from requests import Session

from unshackle.core import __version__
from unshackle.core.vault import Vault


class InsertResult(Enum):
    FAILURE = 0
    SUCCESS = 1
    ALREADY_EXISTS = 2


class HTTPAPI(Vault):
    """Key Vault using a structured HTTP API with JSON payloads and token authentication."""

    def __init__(self, name: str, host: str, password: str):
        super().__init__(name)
        self.url = host
        self.password = password  # This is the API token
        self.current_title = None
        self.session = Session()
        self.session.headers.update({"User-Agent": f"Devine v{__version__}"})
        self.api_session_id = None

    def request(self, method: str, params: dict = None) -> dict:
        """Make a request to the HTTPAPI vault."""
        request_payload = {
            "method": method,
            "params": {
                **(params or {}),
                "session_id": self.api_session_id,
            },
            "token": self.password,
        }

        r = self.session.post(self.url, json=request_payload)

        if r.status_code == 404:
            return {"status": "not_found"}

        if not r.ok:
            raise ValueError(f"API returned HTTP Error {r.status_code}: {r.reason.title()}")

        try:
            res = r.json()
        except json.JSONDecodeError:
            if r.status_code == 404:
                return {"status": "not_found"}
            raise ValueError(f"API returned an invalid response: {r.text}")

        if res.get("status_code") != 200:
            raise ValueError(f"API returned an error: {res['status_code']} - {res['message']}")

        if session_id := res.get("message", {}).get("session_id"):
            self.api_session_id = session_id

        return res.get("message", res)

    def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
        if isinstance(kid, UUID):
            kid = kid.hex

        try:
            title = getattr(self, "current_title", None)
            response = self.request(
                "GetKey",
                {
                    "kid": kid,
                    "service": service.lower(),
                    "title": title,
                },
            )
            if response.get("status") == "not_found":
                return None
            keys = response.get("keys", [])
            for key_entry in keys:
                if key_entry["kid"] == kid:
                    return key_entry["key"]
        except Exception as e:
            print(f"Failed to get key ({e.__class__.__name__}: {e})")
            return None

        return None

    def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
        """Get all keys for a service - this may need to be implemented based on your API."""
        try:
            response = self.request(
                "GetKeys",
                {
                    "service": service.lower(),
                },
            )
            keys = response.get("keys", [])
            for key_entry in keys:
                yield key_entry["kid"], key_entry["key"]
        except Exception:
            return iter([])

    def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
        if not key or key.count("0") == len(key):
            raise ValueError("You cannot add a NULL Content Key to a Vault.")

        if isinstance(kid, UUID):
            kid = kid.hex

        title = getattr(self, "current_title", None)

        try:
            response = self.request(
                "InsertKey",
                {
                    "kid": kid,
                    "key": key,
                    "service": service.lower(),
                    "title": title,
                },
            )
            if response.get("status") == "not_found":
                return False
            return response.get("inserted", False)
        except Exception:
            return False

    def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
        for kid, key in kid_keys.items():
            if not key or key.count("0") == len(key):
                raise ValueError("You cannot add a NULL Content Key to a Vault.")

        processed_kid_keys = {
            str(kid).replace("-", "") if isinstance(kid, UUID) else kid: key for kid, key in kid_keys.items()
        }

        inserted_count = 0
        title = getattr(self, "current_title", None)

        for kid, key in processed_kid_keys.items():
            try:
                response = self.request(
                    "InsertKey",
                    {
                        "kid": kid,
                        "key": key,
                        "service": service.lower(),
                        "title": title,
                    },
                )
                if response.get("status") == "not_found":
                    continue
                if response.get("inserted", False):
                    inserted_count += 1
            except Exception:
                continue

        return inserted_count

    def get_services(self) -> Iterator[str]:
        """Get available services - this may need to be implemented based on your API."""
        try:
            response = self.request("GetServices")
            services = response.get("services", [])
            for service in services:
                yield service
        except Exception:
            return iter([])

    def set_title(self, title: str):
        """
        Set a title to be used for the next key insertions.
        This is optional and will be sent with add_key requests if available.
        """
        self.current_title = title

    def insert_key_with_result(
        self, service: str, kid: Union[UUID, str], key: str, title: Optional[str] = None
    ) -> InsertResult:
        """
        Insert a key and return detailed result information.
        This method provides more granular feedback than the standard add_key method.
        """
        if not key or key.count("0") == len(key):
            raise ValueError("You cannot add a NULL Content Key to a Vault.")

        if isinstance(kid, UUID):
            kid = kid.hex

        if title is None:
            title = getattr(self, "current_title", None)

        try:
            response = self.request(
                "InsertKey",
                {
                    "kid": kid,
                    "key": key,
                    "service": service.lower(),
                    "title": title,
                },
            )

            if response.get("status") == "not_found":
                return InsertResult.FAILURE

            if response.get("inserted", False):
                return InsertResult.SUCCESS
            else:
                return InsertResult.ALREADY_EXISTS

        except Exception:
            return InsertResult.FAILURE