#!/usr/bin/env python3
"""
vnc-token.py — terbitkan JWT untuk akses noVNC (websockify JWTTokenApi).

SAMA PERSIS dengan yang dilakukan websockify: membaca secret file lalu
import_key(k=<isi file>, kty='oct'), claims {host, port, exp}.

Pemakaian (pakai venv python yang punya jwcrypto):
  /opt/data/home/vncsetup/venv/bin/python vnc-token.py                 # token valid 1 jam
  /opt/data/home/vncsetup/venv/bin/python vnc-token.py --ttl 300
  /opt/data/home/vncsetup/venv/bin/python vnc-token.py --expired       # token kedaluwarsa (tes)

Referensi implementasi Go untuk BE (PENTING — jwcrypto pakai kunci base64url, bukan raw hex):
  import "encoding/base64"
  secretBytes, _ := base64.RawURLEncoding.DecodeString(vncJWTSecret)  // 48 byte
  claims := jwt.MapClaims{"host": "localhost", "port": 5901, "exp": time.Now().Add(time.Hour).Unix()}
  s, _ := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString(secretBytes)
"""
import argparse
import json
import os
import time
from pathlib import Path

from jwcrypto import jwt, jwk

DEFAULT_SECRET_FILE = "/opt/data/home/vncsetup/run/vnc-jwt-secret"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ttl", type=int, default=3600, help="masa berlaku detik (default 3600)")
    ap.add_argument("--expired", action="store_true", help="token yang sudah kedaluwarsa (untuk tes)")
    ap.add_argument("--host", default="localhost")
    ap.add_argument("--port", type=int, default=5901)
    ap.add_argument("--secret-file", default=DEFAULT_SECRET_FILE)
    args = ap.parse_args()

    secret = os.environ.get("VNC_JWT_SECRET") or Path(args.secret_file).read_text().strip()
    key = jwk.JWK()
    # Mirror websockify JWTTokenApi.lookup(): import_key(k=<raw string>, kty='oct')
    try:
        key.import_key(k=secret, kty="oct")
    except Exception as e:
        raise SystemExit(f"ERROR: secret tidak bisa dipakai sebagai oct key: {e}")

    now = int(time.time())
    exp = now - 60 if args.expired else now + args.ttl
    claims = {"host": args.host, "port": args.port, "exp": exp, "iat": now}
    token = jwt.JWT(header={"alg": "HS256", "typ": "JWT"}, claims=json.dumps(claims))
    token.make_signed_token(key)
    print(token.serialize())


if __name__ == "__main__":
    main()