scripts/wmf-update-ssh-config - operations/debs/wmf-sre-laptop

Gitiles

blob: 0212b660a629e9be94237181a3d2e0e9ad053fe9 [file] [log] [blame]

#!/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()

txt

json