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