import base64
import json
import os
import secrets
import sys
import urllib.error
import urllib.request
import uuid
from cryptography.hazmat.primitives.asymmetric import x25519

action = sys.argv[1] if len(sys.argv) > 1 else "ensure"
node = sys.argv[2] if len(sys.argv) > 2 else ""

if not node:
    sys.stderr.write("Node name required\n")
    sys.exit(1)

vault_addr = os.environ.get("VAULT_ADDR", "http://10.0.10.12:8200").rstrip("/")
vault_token = os.environ.get("VAULT_TOKEN", "")

if not vault_token:
    sys.stderr.write("VAULT_TOKEN not found in environment\n")
    sys.exit(1)

endpoint = f"{vault_addr}/v1/secret/data/kontur/{node}"

existing_data = None
req_get = urllib.request.Request(endpoint, headers={"X-Vault-Token": vault_token})
try:
    with urllib.request.urlopen(req_get) as resp:
        body = json.loads(resp.read().decode())
        existing_data = body.get("data", {}).get("data", {})
except urllib.error.HTTPError as e:
    if e.code != 404:
        sys.stderr.write(f"Vault error: {e.code} {e.reason}\n")
        sys.exit(1)

if action == "show":
    if not existing_data:
        sys.stderr.write(f"No secrets found for {node}\n")
        sys.exit(1)
    public_view = {
        "vless_uuid": existing_data.get("vless_uuid"),
        "reality_public_key": existing_data.get("reality_public_key"),
        "reality_short_id": existing_data.get("reality_short_id"),
        "reality_server_name": existing_data.get("reality_server_name"),
        "awg_server_public_key": existing_data.get("awg_server_public_key"),
        "awg_router_public_key": existing_data.get("awg_router_public_key"),
        "awg_h1": existing_data.get("awg_h1"),
        "awg_h2": existing_data.get("awg_h2"),
        "awg_h3": existing_data.get("awg_h3"),
        "awg_h4": existing_data.get("awg_h4")
    }
    sys.stdout.write(json.dumps(public_view, indent=2) + "\n")
    sys.exit(0)

if action == "ensure" and existing_data:
    sys.stdout.write(f"Secrets for {node} already exist in Vault (idempotent, skipped)\n")
    sys.exit(0)

v_srv = x25519.X25519PrivateKey.generate()
v_priv = base64.urlsafe_b64encode(v_srv.private_bytes_raw()).decode().rstrip("=")
v_pub = base64.urlsafe_b64encode(v_srv.public_key().public_bytes_raw()).decode().rstrip("=")

awg_srv = x25519.X25519PrivateKey.generate()
awg_s_priv = base64.b64encode(awg_srv.private_bytes_raw()).decode()
awg_s_pub = base64.b64encode(awg_srv.public_key().public_bytes_raw()).decode()

awg_rtr = x25519.X25519PrivateKey.generate()
awg_r_priv = base64.b64encode(awg_rtr.private_bytes_raw()).decode()
awg_r_pub = base64.b64encode(awg_rtr.public_key().public_bytes_raw()).decode()

payload = {
    "data": {
        "vless_uuid": str(uuid.uuid4()),
        "reality_private_key": v_priv,
        "reality_public_key": v_pub,
        "reality_short_id": secrets.token_hex(4),
        "reality_server_name": "sss.svoo.tech",
        "awg_server_private_key": awg_s_priv,
        "awg_server_public_key": awg_s_pub,
        "awg_router_private_key": awg_r_priv,
        "awg_router_public_key": awg_r_pub,
        "awg_jc": 4,
        "awg_jmin": 40,
        "awg_jmax": 70,
        "awg_s1": 15,
        "awg_s2": 25,
        "awg_h1": secrets.randbelow(4000000000) + 100000000,
        "awg_h2": secrets.randbelow(4000000000) + 100000000,
        "awg_h3": secrets.randbelow(4000000000) + 100000000,
        "awg_h4": secrets.randbelow(4000000000) + 100000000
    }
}

req_post = urllib.request.Request(
    endpoint,
    data=json.dumps(payload).encode(),
    headers={"X-Vault-Token": vault_token, "Content-Type": "application/json"}
)

with urllib.request.urlopen(req_post) as resp:
    if resp.status in (200, 204):
        sys.stdout.write(f"Successfully initialized and stored secrets for {node} in Vault\n")