#!/usr/bin/env python3
# Copyright (C) 2026 Douglas Quigg (dstroy0) <dquigg123@gmail.com>
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# Generate a throwaway SSH host key pair for this server - never committed, made
# on demand. It replaces keeping a private key in the repo: run it to provision a
# board (runtime, via NVS) or to embed a key at compile time.
#
# Two key types:
#   rsa       RSA-2048, stored in NVS as PKCS#8 DER under ssh_host_key/priv_der
#             (the runtime "service" path the firmware loads at boot).
#   ed25519   32-byte seed, easiest to embed at compile time and install with
#             ssh_hostkey_ed25519_set(seed).
#
# Outputs (into --out-dir, default the current directory):
#   <name>.pem / <name>.der    the private key (PEM and DER / raw seed)
#   <name>.pub                 the OpenSSH public key for a client's known_hosts
#   with --header/--symbol: a C header with the key bytes for compile-time embed
#     (RSA -> the PKCS#8 DER; Ed25519 -> the 32-byte seed) + a <SYMBOL>_LEN.
#
# Needs `openssl` and `ssh-keygen` on PATH - the same tools docs/SSH.md uses.
#
# Examples:
#   python3 tools/gen_ssh_host_key.py --type rsa --out-dir ./keys
#   python3 tools/gen_ssh_host_key.py --type ed25519 --header host_key.h --symbol HOST_KEY_SEED

import argparse
import os
import subprocess
import sys


def _run(cmd, **kw):
    return subprocess.run(cmd, check=True, **kw)


def _need(tool):
    from shutil import which

    if which(tool) is None:
        sys.exit("error: %s not found on PATH (needed to generate the key)" % tool)


def _write_key(pem_path, gen_cmd):
    with open(pem_path, "wb") as f:
        _run(gen_cmd, stdout=f, stderr=subprocess.DEVNULL)
    os.chmod(pem_path, 0o600)  # ssh-keygen refuses to read a group/other-accessible key file


def gen_rsa(pem_path):
    _write_key(pem_path, ["openssl", "genrsa", "2048"])
    # PKCS#8 DER is what the firmware stores in NVS (ssh_host_key/priv_der).
    return _run(
        ["openssl", "pkcs8", "-topk8", "-nocrypt", "-in", pem_path, "-outform", "DER"], capture_output=True
    ).stdout


def gen_ed25519(pem_path):
    _write_key(pem_path, ["openssl", "genpkey", "-algorithm", "ed25519"])
    # The PKCS#8 private DER ends with OCTET STRING(0x04 0x20 || seed): take the last 32 bytes.
    priv_der = _run(["openssl", "pkey", "-in", pem_path, "-outform", "DER"], capture_output=True).stdout
    seed = priv_der[-32:]
    if len(seed) != 32:
        sys.exit("error: could not extract the 32-byte Ed25519 seed from openssl output")
    return seed


def write_pub(pem_path, pub_path):
    # ssh-keygen -y derives the OpenSSH public key from the private key (both types).
    pub = _run(["ssh-keygen", "-y", "-f", pem_path], capture_output=True).stdout
    with open(pub_path, "wb") as f:
        f.write(pub)


def write_header(path, symbol, blob, what):
    lines = [
        "// Generated by tools/gen_ssh_host_key.py - a throwaway SSH host key, DO NOT COMMIT.",
        "// %s" % what,
        "#pragma once",
        "#include <stdint.h>",
        "",
        "static const uint8_t %s[] = {" % symbol,
    ]
    for i in range(0, len(blob), 16):
        lines.append("    " + ", ".join("0x%02x" % b for b in blob[i : i + 16]) + ",")
    lines.append("};")
    lines.append("static const unsigned int %s_LEN = %d;" % (symbol, len(blob)))
    with open(path, "w") as f:
        f.write("\n".join(lines) + "\n")


def main():
    ap = argparse.ArgumentParser(description="Generate a throwaway SSH host key pair.")
    ap.add_argument("--type", choices=["rsa", "ed25519"], default="rsa")
    ap.add_argument("--out-dir", default=".")
    ap.add_argument("--name", default="ssh_host_key")
    ap.add_argument("--header", help="also emit a C header with the key bytes (compile-time embed)")
    ap.add_argument("--symbol", default="SSH_HOST_KEY", help="C symbol name for --header")
    args = ap.parse_args()

    _need("openssl")
    _need("ssh-keygen")
    os.makedirs(args.out_dir, exist_ok=True)
    pem = os.path.join(args.out_dir, args.name + ".pem")
    der_path = os.path.join(args.out_dir, args.name + ".der")
    pub = os.path.join(args.out_dir, args.name + ".pub")

    if args.type == "rsa":
        blob = gen_rsa(pem)  # PKCS#8 DER
        header_what = "RSA-2048 host key, PKCS#8 DER (NVS ssh_host_key/priv_der)."
    else:
        blob = gen_ed25519(pem)  # 32-byte seed
        header_what = "Ed25519 host key, 32-byte seed for ssh_hostkey_ed25519_set()."
    with open(der_path, "wb") as f:
        f.write(blob)
    write_pub(pem, pub)
    if args.header:
        write_header(args.header, args.symbol, blob, header_what)

    print("generated %s host key (%d-byte %s):" % (args.type, len(blob), "seed" if args.type == "ed25519" else "DER"))
    print("  private : %s / %s" % (pem, der_path))
    print("  public  : %s   (append to a client's ~/.ssh/known_hosts)" % pub)
    if args.header:
        print("  header  : %s   (symbol %s, embed at compile time)" % (args.header, args.symbol))
    print("\nKeep the private key secret; it is not committed. See docs/SSH.md 'Host key provisioning'.")


if __name__ == "__main__":
    main()
