#!/usr/bin/env python3
"""
Подпись паспорта ИИ-агента (формат 0.1) — уровень К3.

Ключ Ed25519, подпись JWS с отделённым содержимым, как у карточки агента A2A:
  signatures[].protected = base64url({"alg":"EdDSA","kid":"..."})
  подписывается строка  protected + "." + base64url(JCS(паспорт без поля signatures))
  JCS — каноническая форма JSON по RFC 8785.

Закрытый ключ создаётся и хранится только у вас. На сайт кладётся паспорт
с открытым ключом в keys[] и подписью в signatures[].

Нужно: Python 3.8+ и пакет cryptography  (pip install cryptography)

Команды:
  python3 sign-passport.py keygen --kid key-2026
      создаёт agent-passport-private.pem (храните в секрете) и печатает открытый ключ
  python3 sign-passport.py sign agent-passport.json --key agent-passport-private.pem --kid key-2026
      добавляет открытый ключ в keys[], подписывает, сохраняет agent-passport.signed.json
  python3 sign-passport.py verify agent-passport.signed.json
      проверяет подписи локально

Источник: https://паспорт-агента.рф/vypusk/#podpis   Лицензия: CC BY 4.0
"""
import argparse
import base64
import json
import os
import sys

try:
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
    from cryptography.hazmat.primitives import serialization
    from cryptography.exceptions import InvalidSignature
except ImportError:
    sys.exit("Нужен пакет cryptography: pip install cryptography")


def b64u(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode("ascii")


def b64u_dec(s: str) -> bytes:
    return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def jcs(v) -> str:
    """Каноническая форма JSON (RFC 8785)."""
    if isinstance(v, dict):
        keys = sorted(v.keys(), key=lambda k: k.encode("utf-16-be"))
        return "{" + ",".join(json.dumps(k, ensure_ascii=False) + ":" + jcs(v[k]) for k in keys) + "}"
    if isinstance(v, list):
        return "[" + ",".join(jcs(x) for x in v) + "]"
    if isinstance(v, bool):
        return "true" if v else "false"
    if v is None:
        return "null"
    if isinstance(v, int):
        return str(v)
    if isinstance(v, float):
        if v.is_integer() and abs(v) < 1e15:
            return str(int(v))
        return repr(v)
    return json.dumps(v, ensure_ascii=False)


def passports_of(doc):
    if isinstance(doc, dict) and "passports" in doc:
        return doc["passports"]
    return [doc]


def load_private(path):
    with open(path, "rb") as f:
        return serialization.load_pem_private_key(f.read(), password=None)


def public_x(priv) -> str:
    raw = priv.public_key().public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    return b64u(raw)


def cmd_keygen(a):
    path = a.out or "agent-passport-private.pem"
    if os.path.exists(path):
        sys.exit(f"Файл {path} уже есть — не перезаписываю. Укажите другое имя через --out")
    priv = Ed25519PrivateKey.generate()
    pem = priv.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, "wb") as f:
        f.write(pem)
    jwk = {"kid": a.kid, "kty": "OKP", "crv": "Ed25519", "x": public_x(priv), "use": "sig"}
    print(f"Закрытый ключ: {path} (храните в секрете, на сайт не кладите)")
    print("Открытый ключ для keys[]:")
    print(json.dumps(jwk, ensure_ascii=False, indent=2))


def cmd_sign(a):
    priv = load_private(a.key)
    with open(a.passport, encoding="utf-8") as f:
        doc = json.load(f)
    jwk = {"kid": a.kid, "kty": "OKP", "crv": "Ed25519", "x": public_x(priv), "use": "sig"}
    protected = b64u(json.dumps({"alg": "EdDSA", "kid": a.kid}, separators=(",", ":")).encode())
    for p in passports_of(doc):
        keys = [k for k in p.get("keys", []) if k.get("kid") != a.kid]
        keys.append(jwk)
        p["keys"] = keys
        p.pop("signatures", None)
        payload = b64u(jcs(p).encode("utf-8"))
        sig = priv.sign((protected + "." + payload).encode("ascii"))
        p["signatures"] = [{"protected": protected, "signature": b64u(sig)}]
    out = a.out or a.passport.rsplit(".json", 1)[0] + ".signed.json"
    with open(out, "w", encoding="utf-8") as f:
        json.dump(doc, f, ensure_ascii=False, indent=2)
        f.write("\n")
    print(f"Готово: {out}")
    print("Положите его на сайт по адресу /.well-known/agent-passport.json и проверьте на https://паспорт-агента.рф/proverka/")


def cmd_verify(a):
    with open(a.passport, encoding="utf-8") as f:
        doc = json.load(f)
    bad = 0
    for i, p in enumerate(passports_of(doc)):
        keys = {k.get("kid"): k for k in p.get("keys", [])}
        body = {k: v for k, v in p.items() if k != "signatures"}
        payload = b64u(jcs(body).encode("utf-8"))
        sigs = p.get("signatures", [])
        if not sigs:
            print(f"[{i}] подписи нет — уровень К2")
            continue
        for s in sigs:
            hdr = json.loads(b64u_dec(s["protected"]))
            k = keys.get(hdr.get("kid"))
            if hdr.get("alg") != "EdDSA" or not k:
                print(f"[{i}] ключ {hdr.get('kid')} не найден или алгоритм не EdDSA")
                bad += 1
                continue
            try:
                Ed25519PublicKey.from_public_bytes(b64u_dec(k["x"])).verify(b64u_dec(s["signature"]), (s["protected"] + "." + payload).encode("ascii"))
                print(f"[{i}] подпись верна (ключ {hdr['kid']})")
            except InvalidSignature:
                print(f"[{i}] подпись НЕ сходится")
                bad += 1
    sys.exit(1 if bad else 0)


def main():
    ap = argparse.ArgumentParser(description="Подпись паспорта ИИ-агента (формат 0.1)")
    sub = ap.add_subparsers(dest="cmd", required=True)
    k = sub.add_parser("keygen", help="создать пару ключей Ed25519")
    k.add_argument("--kid", required=True, help="имя ключа, например key-2026")
    k.add_argument("--out", help="файл закрытого ключа")
    k.set_defaults(fn=cmd_keygen)
    s = sub.add_parser("sign", help="подписать паспорт")
    s.add_argument("passport")
    s.add_argument("--key", required=True)
    s.add_argument("--kid", required=True)
    s.add_argument("--out")
    s.set_defaults(fn=cmd_sign)
    v = sub.add_parser("verify", help="проверить подписи")
    v.add_argument("passport")
    v.set_defaults(fn=cmd_verify)
    a = ap.parse_args()
    a.fn(a)


if __name__ == "__main__":
    main()