Giuseppe Lavagetto | 612c942 | 2020-11-08 08:55:22 +0100 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
| 2 | import errno |
| 3 | import getpass |
| 4 | import os |
| 5 | import subprocess |
| 6 | from pathlib import Path |
| 7 | from typing import Dict |
| 8 | |
| 9 | ENVS = ("prod", "cloud") |
| 10 | |
David Caro | 406635d | 2021-03-02 17:33:28 +0100 | [diff] [blame^] | 11 | CONF_MARKER = "### wmf-laptop-sre - do not edit below this point ###\n" |
Giuseppe Lavagetto | 612c942 | 2020-11-08 08:55:22 +0100 | [diff] [blame] | 12 | |
| 13 | |
| 14 | def check_key(ssh_dir: Path, env: str) -> str: |
| 15 | """ |
| 16 | Checks if a key with the standard name exists. |
| 17 | |
| 18 | Returns the full path to the key""" |
| 19 | keyname = ssh_dir.joinpath(f"id_wmf_{env}") |
| 20 | # Now if the key exists, we'll just return its value |
| 21 | if keyname.is_file(): |
| 22 | return str(keyname) |
| 23 | # Else, raise an error. |
| 24 | raise ValueError(f"ssh key {keyname} not found.") |
| 25 | |
| 26 | |
| 27 | def update_ssh_config(ssh_dir: Path, keys: Dict): |
| 28 | """Update the WMF section of the ssh configuration""" |
David Caro | 406635d | 2021-03-02 17:33:28 +0100 | [diff] [blame^] | 29 | with open("/usr/share/wmf-laptop-sre/ssh-client-config", "r") as fh: |
Giuseppe Lavagetto | 612c942 | 2020-11-08 08:55:22 +0100 | [diff] [blame] | 30 | ssh_config_tpl = fh.read() |
| 31 | # Now let's get the variables to substitute. |
| 32 | user = input("Please provide your production/cloud shell username: ") |
| 33 | new_config = CONF_MARKER |
| 34 | new_config += ( |
| 35 | ssh_config_tpl.replace("USERNAME", user) |
| 36 | .replace("PRODUCTION_KEY", os.path.basename(keys["prod"])) |
| 37 | .replace("WMCS_KEY", os.path.basename(keys["cloud"])) |
| 38 | ) |
| 39 | # Now let's replace the old config. |
| 40 | configfile = ssh_dir.joinpath("config") |
| 41 | config = "" |
| 42 | if configfile.is_file(): |
| 43 | print("Patching the ssh configuration") |
| 44 | with configfile.open("r") as fh: |
| 45 | for line in fh: |
| 46 | if line == CONF_MARKER: |
| 47 | break |
| 48 | config += line |
| 49 | config += new_config |
| 50 | with configfile.open("w") as fh: |
| 51 | fh.write(config) |
| 52 | |
| 53 | |
| 54 | def start_systemd(env: str): |
| 55 | """Enable the systemd service if needed""" |
| 56 | svcname = f"ssh-agent@wmf-{env}.service" |
| 57 | try: |
| 58 | subprocess.check_call(f"systemctl --user is-active {svcname} > /dev/null", shell=True) |
| 59 | except subprocess.CalledProcessError: |
| 60 | print(f"Enabling {svcname}") |
| 61 | subprocess.check_call(f"systemctl --user start {svcname}", shell=True) |
| 62 | |
| 63 | def make_known_hosts_dir(ssh_dir: Path): |
| 64 | """Create the known_hosts directory if not present""" |
| 65 | kh = ssh_dir.joinpath("known_hosts.d") |
| 66 | print(f"Creating {kh} if non-existent") |
| 67 | kh.mkdir(parents=True, exist_ok=True) |
| 68 | |
| 69 | def main(): |
| 70 | ssh_dir = Path.home().joinpath(".ssh") |
| 71 | keys = {} |
| 72 | for env in ENVS: |
| 73 | keyname = check_key(ssh_dir, env) |
| 74 | keys[env] = keyname |
| 75 | # Now check if the user already has an ssh config, |
| 76 | # In that case just paste the wmf config below it. |
| 77 | update_ssh_config(ssh_dir, keys) |
| 78 | make_known_hosts_dir(ssh_dir) |
| 79 | # Now let's ensure the services are up and running. |
| 80 | for env in ENVS: |
| 81 | start_systemd(env) |
| 82 | |
| 83 | |
| 84 | if __name__ == "__main__": |
| 85 | main() |