pikasTech-unidesk/scripts/src/deploy-ssh-identity.ts

import { spawnSync } from "node:child_process";
import { createHash } from "node:crypto";
import { existsSync, readFileSync } from "node:fs";
import { type UniDeskConfig, repoRoot, rootPath } from "./config";

export interface GithubSshIdentityDistribution {
  ok: boolean;
  detail: string;
  fingerprint: string | null;
  seededFromLocal: boolean;
  raw?: unknown;
}

interface GithubSshIdentity {
  privateKey: string;
  publicKey: string;
  knownHosts: string;
  updatedAt: string | null;
}

const identityId = "github.com";
const defaultPrivateKeyPath = "/root/.ssh/id_ed25519";
const defaultKnownHostsPath = "/root/.ssh/known_hosts";

function pgLiteral(value: string): string {
  return `'${value.replace(/'/gu, "''")}'`;
}

function commandOutput(command: string[], input?: string, timeoutMs = 30_000): { ok: boolean; stdout: string; stderr: string; exitCode: number | null } {
  const result = spawnSync(command[0], command.slice(1), {
    cwd: repoRoot,
    input,
    encoding: "utf8",
    maxBuffer: 1024 * 1024 * 8,
    timeout: timeoutMs,
  });
  return {
    ok: result.status === 0,
    stdout: result.stdout ?? "",
    stderr: result.stderr ?? result.error?.message ?? "",
    exitCode: result.status,
  };
}

function runPsql(config: UniDeskConfig, sql: string): { ok: boolean; stdout: string; stderr: string; exitCode: number | null } {
  return commandOutput([
    "docker",
    "exec",
    "-i",
    "unidesk-database",
    "psql",
    "-v",
    "ON_ERROR_STOP=1",
    "-U",
    config.database.user,
    "-d",
    config.database.name,
    "-X",
    "-q",
    "-t",
    "-A",
  ], sql);
}

function ensureIdentityTableSql(): string {
  return `
CREATE TABLE IF NOT EXISTS unidesk_deploy_ssh_identities (
  id TEXT PRIMARY KEY,
  host TEXT NOT NULL,
  private_key TEXT NOT NULL,
  public_key TEXT NOT NULL,
  public_key_fingerprint TEXT NOT NULL DEFAULT '',
  known_hosts TEXT NOT NULL DEFAULT '',
  source TEXT NOT NULL DEFAULT 'operator-local',
  created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
  updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
`;
}

function publicKeyFingerprint(publicKey: string): string {
  const parts = publicKey.trim().split(/\s+/u);
  const encoded = parts[1] ?? "";
  if (encoded.length === 0) return "";
  const digest = createHash("sha256").update(Buffer.from(encoded, "base64")).digest("base64").replace(/=+$/u, "");
  return `SHA256:${digest}`;
}

function githubKnownHostLinesFromText(text: string): string[] {
  const rows: string[] = [];
  const seen = new Set<string>();
  for (const rawLine of text.split(/\r?\n/u)) {
    const line = rawLine.trim();
    if (line.length === 0 || line.startsWith("#")) continue;
    const parts = line.split(/\s+/u);
    if (parts.length < 3) continue;
    const hostList = parts[0]?.split(",") ?? [];
    if (!hostList.some((host) => host === "github.com" || host === "[github.com]:22")) continue;
    if (seen.has(line)) continue;
    seen.add(line);
    rows.push(line);
  }
  return rows;
}

function readGithubKnownHosts(): string {
  const knownHostsPath = process.env.UNIDESK_GITHUB_KNOWN_HOSTS_PATH || defaultKnownHostsPath;
  const localLines = existsSync(knownHostsPath) ? githubKnownHostLinesFromText(readFileSync(knownHostsPath, "utf8")) : [];
  if (localLines.length > 0) return `${localLines.join("\n")}\n`;
  const scan = commandOutput(["ssh-keyscan", "-t", "rsa,ecdsa,ed25519", "github.com"], undefined, 20_000);
  const scannedLines = githubKnownHostLinesFromText(scan.stdout);
  if (scannedLines.length === 0) throw new Error(scan.stderr || "failed to collect github.com SSH host keys");
  return `${scannedLines.join("\n")}\n`;
}

function readLocalGithubIdentity(): GithubSshIdentity | null {
  const privateKeyPath = process.env.UNIDESK_GITHUB_SSH_KEY_PATH || defaultPrivateKeyPath;
  if (!existsSync(privateKeyPath)) return null;
  const privateKey = readFileSync(privateKeyPath, "utf8");
  if (!/-----BEGIN [A-Z ]*PRIVATE KEY-----/u.test(privateKey)) throw new Error(`invalid private key format: ${privateKeyPath}`);
  const publicKeyPath = `${privateKeyPath}.pub`;
  const publicKey = existsSync(publicKeyPath)
    ? readFileSync(publicKeyPath, "utf8").trim()
    : commandOutput(["ssh-keygen", "-y", "-f", privateKeyPath]).stdout.trim();
  if (!/^(ssh-ed25519|ssh-rsa|ecdsa-sha2-nistp256)\s+\S+/u.test(publicKey)) throw new Error(`invalid public key for ${privateKeyPath}`);
  return {
    privateKey: privateKey.endsWith("\n") ? privateKey : `${privateKey}\n`,
    publicKey,
    knownHosts: readGithubKnownHosts(),
    updatedAt: null,
  };
}

function upsertGithubIdentity(config: UniDeskConfig, identity: GithubSshIdentity): void {
  const fingerprint = publicKeyFingerprint(identity.publicKey);
  const sql = `
${ensureIdentityTableSql()}
INSERT INTO unidesk_deploy_ssh_identities (id, host, private_key, public_key, public_key_fingerprint, known_hosts, source, updated_at)
VALUES (
  ${pgLiteral(identityId)},
  ${pgLiteral(identityId)},
  ${pgLiteral(identity.privateKey)},
  ${pgLiteral(identity.publicKey)},
  ${pgLiteral(fingerprint)},
  ${pgLiteral(identity.knownHosts)},
  'operator-local',
  now()
)
ON CONFLICT (id) DO UPDATE SET
  host = EXCLUDED.host,
  private_key = EXCLUDED.private_key,
  public_key = EXCLUDED.public_key,
  public_key_fingerprint = EXCLUDED.public_key_fingerprint,
  known_hosts = EXCLUDED.known_hosts,
  source = EXCLUDED.source,
  updated_at = now();
`;
  const result = runPsql(config, sql);
  if (!result.ok) throw new Error(`failed to upsert GitHub SSH identity in PostgreSQL: ${result.stderr || `exit=${result.exitCode}`}`);
}

function readGithubIdentityFromDatabase(config: UniDeskConfig): GithubSshIdentity {
  const sql = `
${ensureIdentityTableSql()}
SELECT json_build_object(
  'privateKey', private_key,
  'publicKey', public_key,
  'knownHosts', known_hosts,
  'updatedAt', updated_at
)::text
FROM unidesk_deploy_ssh_identities
WHERE id = ${pgLiteral(identityId)}
LIMIT 1;
`;
  const result = runPsql(config, sql);
  if (!result.ok) throw new Error(`failed to read GitHub SSH identity from PostgreSQL: ${result.stderr || `exit=${result.exitCode}`}`);
  const line = result.stdout.trim().split(/\r?\n/u).find((item) => item.trim().startsWith("{"));
  if (line === undefined) throw new Error("GitHub SSH identity is missing in PostgreSQL");
  const parsed = JSON.parse(line) as Partial<GithubSshIdentity>;
  const privateKey = typeof parsed.privateKey === "string" ? parsed.privateKey : "";
  const publicKey = typeof parsed.publicKey === "string" ? parsed.publicKey : "";
  const knownHosts = typeof parsed.knownHosts === "string" ? parsed.knownHosts : "";
  if (!/-----BEGIN [A-Z ]*PRIVATE KEY-----/u.test(privateKey)) throw new Error("PostgreSQL GitHub SSH identity has invalid private key material");
  if (!/^(ssh-ed25519|ssh-rsa|ecdsa-sha2-nistp256)\s+\S+/u.test(publicKey)) throw new Error("PostgreSQL GitHub SSH identity has invalid public key material");
  if (githubKnownHostLinesFromText(knownHosts).length === 0) throw new Error("PostgreSQL GitHub SSH identity has no github.com known_hosts rows");
  return {
    privateKey: privateKey.endsWith("\n") ? privateKey : `${privateKey}\n`,
    publicKey,
    knownHosts: knownHosts.endsWith("\n") ? knownHosts : `${knownHosts}\n`,
    updatedAt: typeof parsed.updatedAt === "string" ? parsed.updatedAt : null,
  };
}

function remoteInstallPythonSource(): string {
  return String.raw`
import json
import os
import pathlib
import stat
import subprocess
import sys

data = json.load(sys.stdin)
ssh_dir = pathlib.Path("/home/ubuntu/.ssh")
private_key = str(data.get("privateKey") or "")
public_key = str(data.get("publicKey") or "").strip()
known_hosts = str(data.get("knownHosts") or "")

if "PRIVATE KEY-----" not in private_key:
    raise SystemExit("invalid private key payload")
if not public_key.startswith(("ssh-ed25519 ", "ssh-rsa ", "ecdsa-sha2-nistp256 ")):
    raise SystemExit("invalid public key payload")

ssh_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
os.chmod(ssh_dir, 0o700)
private_path = ssh_dir / "id_ed25519"
public_path = ssh_dir / "id_ed25519.pub"
known_hosts_path = ssh_dir / "known_hosts"

private_path.write_text(private_key if private_key.endswith("\n") else private_key + "\n")
public_path.write_text(public_key + "\n")
os.chmod(private_path, stat.S_IRUSR | stat.S_IWUSR)
os.chmod(public_path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH)

derived = subprocess.run(["ssh-keygen", "-y", "-f", str(private_path)], text=True, capture_output=True)
if derived.returncode != 0:
    raise SystemExit("installed private key failed ssh-keygen verification")
if derived.stdout.strip() != public_key:
    raise SystemExit("installed private key does not match public key")

known_hosts_path.touch(mode=0o600, exist_ok=True)
subprocess.run(["ssh-keygen", "-R", "github.com", "-f", str(known_hosts_path)], text=True, capture_output=True)
subprocess.run(["ssh-keygen", "-R", "[github.com]:22", "-f", str(known_hosts_path)], text=True, capture_output=True)
existing = known_hosts_path.read_text().splitlines()
rows = []
seen = set()
for line in existing + known_hosts.splitlines():
    line = line.strip()
    if not line or line.startswith("#"):
        continue
    parts = line.split()
    if len(parts) < 3:
        continue
    if line in seen:
        continue
    seen.add(line)
    rows.append(line)
known_hosts_path.write_text("\n".join(rows) + "\n")
os.chmod(known_hosts_path, stat.S_IRUSR | stat.S_IWUSR)
print(f"github_ssh_identity_written path={private_path} known_hosts_rows={len(rows)}")
`;
}

function shellQuote(value: string): string {
  return `'${value.replace(/'/gu, `'\\''`)}'`;
}

function distributeGithubIdentity(providerId: string, identity: GithubSshIdentity): { ok: boolean; detail: string; raw: unknown } {
  const payload = JSON.stringify({
    privateKey: identity.privateKey,
    publicKey: identity.publicKey,
    knownHosts: identity.knownHosts,
  });
  const remotePython = remoteInstallPythonSource();
  const remoteCommand = [
    `python3 -c ${shellQuote(remotePython)}`,
    "auth_output=$(ssh -o BatchMode=yes -o IdentitiesOnly=yes -o StrictHostKeyChecking=yes -T git@github.com 2>&1 || true)",
    "printf '%s\\n' \"$auth_output\"",
    "printf '%s\\n' \"$auth_output\" | grep -q 'successfully authenticated'",
  ].join(" && ");
  const result = spawnSync(process.execPath, [
    rootPath("scripts", "cli.ts"),
    "ssh",
    providerId,
    "--",
    remoteCommand,
  ], {
    cwd: repoRoot,
    input: payload,
    encoding: "utf8",
    timeout: 90_000,
    maxBuffer: 1024 * 1024 * 4,
    env: { ...process.env, UNIDESK_SSH_OPEN_TIMEOUT_MS: process.env.UNIDESK_SSH_OPEN_TIMEOUT_MS || "60000" },
  });
  const stdout = result.stdout ?? "";
  const stderr = result.stderr ?? result.error?.message ?? "";
  return {
    ok: result.status === 0,
    detail: [stdout, stderr].filter(Boolean).join("\n").slice(-2000),
    raw: { exitCode: result.status, stdoutTail: stdout.slice(-1200), stderrTail: stderr.slice(-1200) },
  };
}

export async function ensureGithubSshIdentityForProvider(config: UniDeskConfig, providerId: string): Promise<GithubSshIdentityDistribution> {
  let seededFromLocal = false;
  const localIdentity = readLocalGithubIdentity();
  if (localIdentity !== null) {
    upsertGithubIdentity(config, localIdentity);
    seededFromLocal = true;
  }
  const identity = readGithubIdentityFromDatabase(config);
  const fingerprint = publicKeyFingerprint(identity.publicKey);
  const distribution = distributeGithubIdentity(providerId, identity);
  if (!distribution.ok) {
    return {
      ok: false,
      detail: `failed to distribute GitHub SSH identity from PostgreSQL to ${providerId}: ${distribution.detail || "remote ssh failed"}`,
      fingerprint,
      seededFromLocal,
      raw: distribution.raw,
    };
  }
  return {
    ok: true,
    detail: `GitHub SSH identity distributed from PostgreSQL to ${providerId}; fingerprint=${fingerprint}; seededFromLocal=${seededFromLocal}`,
    fingerprint,
    seededFromLocal,
    raw: distribution.raw,
  };
}