api.py

# ═══════════════════════════════════════════════════════════════════
#  SENTINEL GATEWAY — FastAPI External Agent Interface
#  Version 2.0 — Shared persistent Token Authority with key rotation
#
#  Runs separately from the Streamlit UI:
#    uvicorn api:app --host 0.0.0.0 --port 8000
#
#  Both this process and sentinel_gateway.py share the same signing
#  key stored in the signing_key table. Tokens issued by either
#  process are verifiable by the other.
# ═══════════════════════════════════════════════════════════════════

import json, hashlib, time, uuid
from datetime import datetime

try:
    import db_compat as sqlite3          # PostgreSQL adapter (Replit)
except ImportError:
    import sqlite3                       # fallback for local SQLite

try:
    from fastapi import FastAPI, HTTPException, Depends
    from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
    from pydantic import BaseModel
except ImportError:
    raise RuntimeError("Run: pip install fastapi uvicorn")

try:
    import nacl.signing, nacl.exceptions
except ImportError:
    raise RuntimeError("Run: pip install PyNaCl")

import os
DB_PATH = os.environ.get("DATABASE_URL", "sentinel_gateway.db")

GRACE_SECONDS  = 3600  # previous key stays valid this long after rotation
MEMORY_LIMITS  = {"short_term": 3000, "long_term": 10000}

# ── Database connection ──────────────────────────────────────────
def get_conn():
    return sqlite3.connect(DB_PATH, check_same_thread=False)

conn = get_conn()


# ── Audit helper ─────────────────────────────────────────────────
class _Audit:
    def log(self, prompt_id, event_type, action, details, status):
        conn.execute(
            "INSERT INTO audit (ts,prompt_id,event_type,action,details,status) "
            "VALUES (?,?,?,?,?,?)",
            (int(time.time()), prompt_id, event_type,
             action, json.dumps(details), status))
        conn.commit()

    def nonce_used(self, nonce):
        return conn.execute(
            "SELECT 1 FROM used_nonces WHERE nonce=?",
            (nonce,)).fetchone() is not None

    def consume_nonce(self, nonce):
        conn.execute("INSERT OR IGNORE INTO used_nonces VALUES (?,?)",
                     (nonce, int(time.time())))
        conn.commit()

    def get_recent(self, limit=100):
        rows = conn.execute(
            "SELECT * FROM audit ORDER BY ts DESC LIMIT ?",
            (limit,)).fetchall()
        return [{"id": r[0],
                 "time": datetime.fromtimestamp(r[1]).strftime("%H:%M:%S"),
                 "prompt_id": r[2], "event": r[3],
                 "action": r[4], "details": r[5], "status": r[6]}
                for r in rows]


# ── Prompt registry ──────────────────────────────────────────────
class _PromptReg:
    def get(self, prompt_id):
        row = conn.execute(
            "SELECT * FROM prompts WHERE prompt_id=?",
            (prompt_id,)).fetchone()
        if not row:
            return None
        return {"prompt_id": row[0], "user_id": row[1],
                "instruction": row[2], "scope": json.loads(row[3]),
                "issued_at": row[4]}

    def register(self, prompt_id, user_id, instruction, scope, token):
        token_hash = hashlib.sha256(
            json.dumps(token, sort_keys=True).encode()).hexdigest()
        conn.execute(
            "INSERT OR REPLACE INTO prompts VALUES (?,?,?,?,?,?)",
            (prompt_id, user_id, instruction,
             json.dumps(scope), int(time.time()), token_hash))
        conn.commit()


# ── Agent registry ───────────────────────────────────────────────
class _AgentReg:
    def verify_api_key(self, api_key):
        key_hash = hashlib.sha256(api_key.encode()).hexdigest()
        row = conn.execute(
            "SELECT * FROM agents WHERE api_key_hash=? AND active=1",
            (key_hash,)).fetchone()
        if not row:
            return None
        return {"agent_id": row[0], "agent_name": row[1],
                "agent_type": row[2], "scope_ceiling": json.loads(row[3])}


# ── Shared, persisted Token Authority ───────────────────────────
class _TokenAuthority:
    """
    Single source of truth for signing keys — shared with the Streamlit
    process via the signing_key table. Both processes read the same key
    on startup; tokens issued by either are verifiable by the other.

    Key rotation:
      - rotate() promotes current → previous_key_bytes
      - previous key remains valid for GRACE_SECONDS (default 1 hour)
      - rotation is logged to the audit table
    """

    def __init__(self, db_conn):
        self._conn        = db_conn
        self.previous_key = None
        self.rotated_at   = 0
        row = db_conn.execute(
            "SELECT key_bytes, previous_key_bytes, rotated_at "
            "FROM signing_key WHERE id=1").fetchone()
        if row:
            self.signing_key  = nacl.signing.SigningKey(bytes.fromhex(row[0]))
            if row[1]:
                self.previous_key = nacl.signing.SigningKey(bytes.fromhex(row[1]))
            self.rotated_at = row[2] or 0
        else:
            self.signing_key = nacl.signing.SigningKey.generate()
            db_conn.execute(
                "INSERT INTO signing_key (id, key_bytes, previous_key_bytes, rotated_at) "
                "VALUES (1,?,?,?)",
                (self.signing_key.encode().hex(), "", 0))
            db_conn.commit()
        self.verify_key = self.signing_key.verify_key

    def rotate(self, acting_agent: str) -> int:
        """Rotate key. Returns grace_until timestamp."""
        old_hex           = self.signing_key.encode().hex()
        self.previous_key = self.signing_key
        self.signing_key  = nacl.signing.SigningKey.generate()
        self.verify_key   = self.signing_key.verify_key
        self.rotated_at   = int(time.time())
        self._conn.execute(
            "UPDATE signing_key SET key_bytes=?, previous_key_bytes=?, rotated_at=? WHERE id=1",
            (self.signing_key.encode().hex(), old_hex, self.rotated_at))
        self._conn.commit()
        grace_until = self.rotated_at + GRACE_SECONDS
        audit.log("system", "key_rotated", "rotate",
                  {"acting_agent": acting_agent,
                   "grace_until": datetime.fromtimestamp(grace_until).strftime(
                       "%Y-%m-%d %H:%M:%S")},
                  "rotated")
        return grace_until

    def issue_token(self, prompt_id, user_id, scope, expires_in=600):
        payload = {
            "prompt_id": prompt_id, "user_id": user_id,
            "scope":     scope,
            "issued_at":  int(time.time()),
            "expires_at": int(time.time()) + expires_in,
            "nonce":      uuid.uuid4().hex,
        }
        raw = json.dumps(payload, sort_keys=True).encode()
        sig = self.signing_key.sign(raw).signature.hex()
        return {"payload": payload, "signature": sig}

    def verify_token(self, token):
        raw = json.dumps(token["payload"], sort_keys=True).encode()
        sig = bytes.fromhex(token["signature"])
        # Try current key
        try:
            self.verify_key.verify(raw, sig)
            if int(time.time()) > token["payload"]["expires_at"]:
                return False, "Token expired"
            return True, "Valid"
        except nacl.exceptions.BadSignatureError:
            pass
        except Exception as e:
            return False, str(e)
        # Fall back to previous key within grace window
        if (self.previous_key and self.rotated_at and
                int(time.time()) < self.rotated_at + GRACE_SECONDS):
            try:
                self.previous_key.verify_key.verify(raw, sig)
                if int(time.time()) > token["payload"]["expires_at"]:
                    return False, "Token expired"
                return True, "Valid (grace period)"
            except nacl.exceptions.BadSignatureError:
                pass
            except Exception as e:
                return False, str(e)
        return False, "Invalid signature"


# ── Memory store ─────────────────────────────────────────────────
class _MemoryStore:
    def get(self, memory_type):
        row = conn.execute(
            "SELECT content, updated_at, updated_by FROM user_memory "
            "WHERE memory_type=?", (memory_type,)).fetchone()
        if not row:
            return {"content": "", "updated_at": None, "updated_by": None,
                    "char_limit": MEMORY_LIMITS.get(memory_type, 3000)}
        return {"content": row[0] or "", "updated_at": row[1],
                "updated_by": row[2],
                "char_limit": MEMORY_LIMITS.get(memory_type, 3000)}

    def set(self, memory_type, content, updated_by):
        limit = MEMORY_LIMITS.get(memory_type, 3000)
        if len(content) > limit:
            return False, f"Content exceeds {limit} character limit ({len(content)} chars)"
        conn.execute(
            "INSERT OR REPLACE INTO user_memory "
            "(memory_type, content, char_limit, updated_at, updated_by) "
            "VALUES (?,?,?,?,?)",
            (memory_type, content, limit, int(time.time()), updated_by))
        conn.commit()
        return True, "ok"

    def status(self):
        result = {}
        for mt in ("short_term", "long_term"):
            row = conn.execute(
                "SELECT content FROM user_memory WHERE memory_type=?",
                (mt,)).fetchone()
            result[mt] = "has_content" if (row and row[0] and row[0].strip()) else "empty"
        return result


# ── Instantiate shared objects ───────────────────────────────────
audit        = _Audit()
prompt_reg   = _PromptReg()
agent_reg    = _AgentReg()
ta           = _TokenAuthority(conn)
memory_store = _MemoryStore()

# ── FastAPI app ──────────────────────────────────────────────────
app    = FastAPI(title="Sentinel Gateway API", version="2.0")
bearer = HTTPBearer()


def get_agent(creds: HTTPAuthorizationCredentials = Depends(bearer)):
    agent = agent_reg.verify_api_key(creds.credentials)
    if not agent:
        raise HTTPException(401, "Invalid or unregistered agent API key")
    return agent


class IssueTokenReq(BaseModel):
    prompt_id: str
    scope: list
    expires_in: int = 600

class SubmitReq(BaseModel):
    instruction: str
    token: dict

class ActionReq(BaseModel):
    prompt_id: str
    action_type: str
    action_params: dict
    token: dict


@app.get("/")
def root():
    return {"service": "Sentinel Gateway", "version": "2.0",
            "status": "running", "docs": "/docs"}


@app.post("/v1/issue_token")
def issue_token(req: IssueTokenReq, agent=Depends(get_agent)):
    ceiling = agent["scope_ceiling"]
    invalid = [s for s in req.scope if s not in ceiling]
    if invalid:
        raise HTTPException(403, f"Scope {invalid} exceeds ceiling {ceiling}")
    token = ta.issue_token(req.prompt_id, agent["agent_id"],
                           req.scope, req.expires_in)
    audit.log(req.prompt_id, "token_issued_external", "issue",
              {"agent": agent["agent_name"], "scope": req.scope}, "issued")
    return {**token, "memory_status": memory_store.status()}


@app.post("/v1/submit_instruction")
def submit_instruction(req: SubmitReq, agent=Depends(get_agent)):
    valid, reason = ta.verify_token(req.token)
    if not valid:
        return {"status": "blocked", "reason": reason}
    pid   = req.token["payload"]["prompt_id"]
    scope = req.token["payload"]["scope"]
    prompt_reg.register(pid, agent["agent_id"], req.instruction, scope, req.token)
    audit.log(pid, "instruction_registered", "submit",
              {"instruction": req.instruction[:120],
               "scope": scope, "agent": agent["agent_name"]}, "accepted")
    return {"status": "accepted", "prompt_id": pid}


@app.post("/v1/request_action")
def request_action(req: ActionReq, agent=Depends(get_agent)):
    valid, reason = ta.verify_token(req.token)
    if not valid:
        return {"status": "blocked", "reason": reason}
    nonce = req.token["payload"]["nonce"]
    if audit.nonce_used(nonce):
        return {"status": "blocked", "reason": "Nonce already used"}
    audit.consume_nonce(nonce)
    scope = req.token["payload"]["scope"]
    if req.action_type not in scope:
        audit.log(req.prompt_id, "action_blocked", req.action_type,
                  {"agent": agent["agent_name"], "reason": "not in scope"}, "blocked")
        return {"status": "blocked",
                "reason": f"'{req.action_type}' not in scope {scope}"}

    PERMISSION_SIGNALS = {"web_write", "schedule_task"}
    if req.action_type in PERMISSION_SIGNALS:
        result = (f"[GATEWAY: {req.action_type} PERMITTED]\n"
                  f"Execute using your own tools.")
    elif req.action_type == "agent_talk":
        a_url  = req.action_params.get("agent_url",  "").strip()
        a_id   = req.action_params.get("agent_id",   "").strip()
        a_name = req.action_params.get("agent_name", "").strip()
        if not any([a_url, a_id, a_name]):
            result = "[agent_talk] Provide at least one of: agent_url, agent_id, or agent_name"
        elif a_url:
            result = (f"[GATEWAY: agent_talk PERMITTED]\n"
                      f"Target URL: {a_url}\n"
                      f"POST {{\"message\": ...}} using your own HTTP client.")
        else:
            lookup_col = "agent_id" if a_id else "agent_name"
            lookup_val = a_id or a_name
            extra_cond = "" if a_id else " AND active=1"
            row = conn.execute(
                f"SELECT agent_id,agent_name,agent_type,scope_ceiling,"
                f"registered_at,active FROM agents "
                f"WHERE {lookup_col}=?{extra_cond}",
                (lookup_val,)).fetchone()
            if not row:
                result = f"[agent_talk] No registered agent found for: '{lookup_val}'"
            else:
                result = json.dumps({
                    "source":        "internal",
                    "agent_id":      row[0],
                    "agent_name":    row[1],
                    "agent_type":    row[2],
                    "scope_ceiling": json.loads(row[3]),
                    "registered_at": datetime.fromtimestamp(row[4]).strftime("%Y-%m-%d %H:%M"),
                    "active":        bool(row[5]),
                }, indent=2)
    else:
        result = f"[PERMITTED] {req.action_type} — execute with your own connection."

    audit.log(req.prompt_id, "action_permitted_external", req.action_type,
              {"agent": agent["agent_name"],
               "params": str(req.action_params)[:100]}, "executed")
    return {"status": "permitted", "action": req.action_type, "result": result}


@app.post("/v1/rotate_key")
def rotate_key(agent=Depends(get_agent)):
    """
    Rotate the signing key. The previous key remains valid for GRACE_SECONDS (1 hour)
    so in-flight tokens are not orphaned. Restrict access to this endpoint at the
    infrastructure level (e.g. nginx allow-list) — only admin agents should call it.
    Rotation is logged to the audit table.
    """
    grace_until = ta.rotate(acting_agent=agent["agent_name"])
    return {
        "status":         "rotated",
        "grace_until":    datetime.fromtimestamp(grace_until).strftime("%Y-%m-%d %H:%M:%S"),
        "grace_seconds":  GRACE_SECONDS,
        "message":        (f"New key active. Previous key valid until "
                           f"{datetime.fromtimestamp(grace_until).strftime('%H:%M:%S')} "
                           f"({GRACE_SECONDS // 60} min grace period)."),
    }


@app.get("/v1/memory")
def read_memory(type: str = "short_term", agent=Depends(get_agent)):
    """Read short-term or long-term memory. No scope restriction — all agents may read."""
    if type not in ("short_term", "long_term"):
        raise HTTPException(400, "type must be 'short_term' or 'long_term'")
    data = memory_store.get(type)
    return {"type": type, **data}


class WriteMemoryReq(BaseModel):
    type: str
    content: str
    token: dict

@app.post("/v1/memory")
def write_memory(req: WriteMemoryReq, agent=Depends(get_agent)):
    """Write short-term or long-term memory. Requires 'memory_write' in token scope."""
    valid, reason = ta.verify_token(req.token)
    if not valid:
        return {"status": "blocked", "reason": reason}
    if req.type not in ("short_term", "long_term"):
        raise HTTPException(400, "type must be 'short_term' or 'long_term'")
    scope = req.token["payload"]["scope"]
    if "memory_write" not in scope:
        pid = req.token["payload"].get("prompt_id", "unknown")
        audit.log(pid, "memory_write_blocked", "memory_write",
                  {"agent": agent["agent_name"], "reason": "not in scope"}, "blocked")
        return {"status": "blocked",
                "reason": "'memory_write' not in token scope"}
    pid = req.token["payload"].get("prompt_id", "unknown")
    ok, msg = memory_store.set(req.type, req.content, updated_by=agent["agent_name"])
    if not ok:
        return {"status": "error", "reason": msg}
    audit.log(pid, "memory_written", "memory_write",
              {"agent": agent["agent_name"], "type": req.type,
               "chars": len(req.content)}, "ok")
    return {"status": "ok", "type": req.type,
            "chars": len(req.content),
            "char_limit": MEMORY_LIMITS[req.type]}


@app.get("/v1/roles")
def get_roles(agent=Depends(get_agent)):
    rows = conn.execute(
        "SELECT role_name, definition FROM agent_roles ORDER BY role_name"
    ).fetchall()
    return [{"role_name": r[0], "definition": r[1]} for r in rows]


@app.get("/v1/audit")
def get_audit(agent=Depends(get_agent)):
    return audit.get_recent(100)


@app.get("/v1/prompt/{prompt_id}")
def get_prompt(prompt_id: str, agent=Depends(get_agent)):
    record = prompt_reg.get(prompt_id)
    if not record:
        raise HTTPException(404, "Prompt ID not found")
    return record