scripts/wmf-update-ssh-config

#!/usr/bin/env python3
import errno
import getpass
import os
import subprocess
from pathlib import Path
from typing import Dict

ENVS = ("prod", "cloud")

CONF_MARKER = "### wmf-laptop-sre - do not edit below this point ###\n"


def check_key(ssh_dir: Path, env: str) -> str:
    """
    Checks if a key with the standard name exists.

    Returns the full path to the key"""
    keyname = ssh_dir.joinpath(f"id_wmf_{env}")
    # Now if the key exists, we'll just return its value
    if keyname.is_file():
        return str(keyname)
    # Else, raise an error.
    raise ValueError(f"ssh key {keyname} not found.")


def update_ssh_config(ssh_dir: Path, keys: Dict):
    """Update the WMF section of the ssh configuration"""
    with open("/usr/share/wmf-laptop-sre/ssh-client-config", "r") as fh:
        ssh_config_tpl = fh.read()
    # Now let's get the variables to substitute.
    user = input("Please provide your production/cloud shell username: ")
    new_config = CONF_MARKER
    new_config += (
        ssh_config_tpl.replace("USERNAME", user)
        .replace("PRODUCTION_KEY", os.path.basename(keys["prod"]))
        .replace("WMCS_KEY", os.path.basename(keys["cloud"]))
    )
    # Now let's replace the old config.
    configfile = ssh_dir.joinpath("config")
    config = ""
    if configfile.is_file():
        print("Patching the ssh configuration")
        with configfile.open("r") as fh:
            for line in fh:
                if line == CONF_MARKER:
                    break
                config += line
    config += new_config
    with configfile.open("w") as fh:
        fh.write(config)


def start_systemd(env: str):
    """Enable the systemd service if needed"""
    svcname = f"ssh-agent@wmf-{env}.service"
    try:
        subprocess.check_call(f"systemctl --user is-active {svcname} > /dev/null", shell=True)
    except subprocess.CalledProcessError:
        print(f"Enabling {svcname}")
        subprocess.check_call(f"systemctl --user start {svcname}", shell=True)

def make_known_hosts_dir(ssh_dir: Path):
    """Create the known_hosts directory if not present"""
    kh = ssh_dir.joinpath("known_hosts.d")
    print(f"Creating {kh} if non-existent")
    kh.mkdir(parents=True, exist_ok=True)

def main():
    ssh_dir = Path.home().joinpath(".ssh")
    keys = {}
    for env in ENVS:
        keyname = check_key(ssh_dir, env)
        keys[env] = keyname
    # Now check if the user already has an ssh config,
    # In that case just paste the wmf config below it.
    update_ssh_config(ssh_dir, keys)
    make_known_hosts_dir(ssh_dir)
    # Now let's ensure the services are up and running.
    for env in ENVS:
        start_systemd(env)


if __name__ == "__main__":
    main()