ts-scripts/single-scripts/ssh.ts

#!/usr/bin/env runts -- --runtime-bun

// reference: https://git.hatter.ink/rust-scripts/scriptbase/src/branch/main/ssh-rs/src/main.rs

const { spawn } = require("node:child_process");
const { parseArgs } = require("node:util");
const fs = require("node:fs");
const os = require("node:os");

// reference: https://bun.com/docs/runtime/color
const GREEN = Bun.color("green", "ansi");
const BLUE = Bun.color("lightblue", "ansi");
const YELLOW = Bun.color("yellow", "ansi");
const RESET = "\x1B[0m";

class SshTsArgs {
    forwardAgent?: boolean;
    proxy?: boolean;
    help?: boolean;
    host?: string;
}

interface SshConfig {
    default_forward_agent?: boolean;
    default_proxy?: boolean;
    default_username?: string;
    profiles: Map<string, SshProfile>;
}

interface SshProfile {
    default_username?: string;
    alias?: string[];
    host: string;
    proxy?: boolean;
    forward_agent?: boolean;
    comment?: string;
}

interface UsernameAndHost {
    username?: string;
    host: string;
}

function printSshConfig(sshConfig: SshConfig) {
    const allProfiles = [];
    let maxProfileNameLength = 0;
    let maxProfileHostLength = 0;
    for (const k in sshConfig.profiles) {
        const sshProfile = sshConfig.profiles[k];
        allProfiles.push(k);
        if (k.length > maxProfileNameLength) {
            maxProfileNameLength = k.length;
        }
        if (sshProfile.host.length > maxProfileHostLength) {
            maxProfileHostLength = sshProfile.host.length;
        }
    }
    console.log(
        `${GREEN}[OK   ]${RESET} Total ${allProfiles.length} server(s):`,
    );
    allProfiles.sort((a, b) => a.localeCompare(b));
    for (let i = 0; i < allProfiles.length; i++) {
        const k = allProfiles[i];
        const sshProfile = sshConfig.profiles[k];
        const features = [];
        if (sshProfile.proxy) features.push("proxy");
        if (sshProfile.forward_agent) features.push("forward_agent");
        console.log(
            ` - ${k}${
                " ".repeat(maxProfileNameLength - k.length)
            } : ${BLUE}${sshProfile.host}${
                " ".repeat(maxProfileHostLength - sshProfile.host.length)
            }${RESET} ${YELLOW}${
                (sshProfile.alias && sshProfile.alias.length == 1)
                    ? "alias  "
                    : "aliases"
            }: [${
                (sshProfile.alias && sshProfile.alias.join(", ")) || ""
            }]${RESET} # ${sshProfile.comment}${
                (features.length > 0) ? (" ;[" + features.join(" ") + "]") : ""
            }`,
        );
    }
}

function loadSshConfig(): SshConfig {
    const configFile = os.homedir() + "/.config/ssh-rs-config.json";
    try {
        const sshConfigText = fs.readFileSync(configFile, "utf8");
        return JSON.parse(sshConfigText);
    } catch (e) {
        console.error(`Load config file: ${configFile} failed.`, e);
    }
}

function matchProfile(sshConfig: SshConfig, host: string): SshProfile {
    let profiles = [];
    for (const k in sshConfig.profiles) {
        const sshProfile = sshConfig.profiles[k];
        if (k === host) {
            profiles.push(sshProfile);
        } else if (sshProfile.alias && sshProfile.alias.includes(host)) {
            profiles.push(sshProfile);
        }
    }
    if (profiles.length === 0) {
        return null;
    } else if (profiles.length > 1) {
        throw new Error("Find multiple profiles");
    }
    return profiles[0];
}

function parseUsernameAndHost(usernameAndHost: string): UsernameAndHost {
    if (!usernameAndHost) {
        throw new Error("Empty username@host");
    }
    const usernameAndHostParts = usernameAndHost.split("@");
    if (usernameAndHostParts.length == 1) {
        return { host: usernameAndHostParts[0] };
    }
    if (usernameAndHostParts.length > 2) {
        throw new Error(`Base username@host: ${usernameAndHost}`);
    }
    return { username: usernameAndHostParts[0], host: usernameAndHostParts[1] };
}

async function main() {
    const sshConfig = loadSshConfig();

    if (process.argv.length <= 2) {
        printSshConfig(sshConfig);
        return;
    }
    const args = process.argv.slice(2);
    const options = {
        "forward-agent": { type: "boolean", short: "f" },
        "proxy": { type: "boolean", short: "p" },
        "help": { type: "boolean", short: "h" },
        "host": { type: "string", short: "H" },
    };

    const { values, positionals } = parseArgs({
        args,
        options,
        allowPositionals: true,
        tokens: true,
    });
    const sshTsArgs = values as SshTsArgs;
    if (sshTsArgs.help) {
        console.log("ssh.ts [-h|--help]");
        console.log(
            "ssh.ts [-f|--forward-agent] [-p|--proxy] [-H|--host] [username@]host",
        );
        console.log("ssh.ts [-f|--forward-agent] [-p|--proxy] [username@]host");
        return;
    }

    sshTsArgs.forwardAgent = values["forward-agent"];
    if (!sshTsArgs.host && positionals && positionals.length > 0) {
        sshTsArgs.host = positionals[0];
    }
    if (!sshTsArgs.host) {
        console.error("[ERROR] --host required");
        return;
    }

    const { username, host } = parseUsernameAndHost(sshTsArgs.host);
    const sshProfile = matchProfile(sshConfig, host);
    if (sshProfile === null) {
        console.error("[ERROR] No ssh profile found.");
        return;
    }

    const sshCommand = "ssh";
    const sshArgs = [];
    const sshUsername = username || sshProfile.default_username ||
        sshConfig.default_username || "root";
    const sshForwardAgent = sshTsArgs.forwardAgent ||
        sshProfile.forward_agent || sshConfig.default_forward_agent || true;
    const sshProxy = sshTsArgs.proxy || sshProfile.proxy ||
        sshConfig.default_proxy || false;
    if (sshForwardAgent) {
        sshArgs.push("-o");
        sshArgs.push("ForwardAgent=yes");
    }
    if (sshProxy) {
        sshArgs.push("-o");
        sshArgs.push('"ProxyCommand=nc -X 5 -x 127.0.0.1:1080 %h %p"');
    }
    sshArgs.push(`${sshUsername}@${sshProfile.host}`);

    console.log(`${GREEN}[OK   ]${RESET} ${sshCommand} ${sshArgs.join(" ")}`);
    spawn(sshCommand, sshArgs, {
        shell: true,
        stdio: ["inherit", "inherit", "inherit"],
    });
}
await main();