gerrit.wikimedia.org
Gitiles
Code ReviewSign In
gerrit.wikimedia.org / operations / debs / wmf-sre-laptop / da48c987fcf1a1b0191074bfe26d0e14dac7fceb / . / scripts / wmf-update-ssh-config
blob: 4ceddddb0676b257a78a17b66adcc5d302c247b4 [file] [log] [blame]
#!/usr/bin/env python3
import argparse
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​(​configfile​: 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.
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():
parser = argparse​.​ArgumentParser​()
parser​.​add_argument​(​"configfile"​, metavar="OUTPUT_FILE", nargs="?", default="",
help="File to write SSH config to (defaults to ~/.ssh/config)")
parser.add_argument("-s", "--skip-systemd",
help="Skip systemd service activation",
action="store_true")
args = parser.parse_args()
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.
if args.configfile != "":
configfile = Path(args.configfile)
else:
configfile = ssh_dir.joinpath("config")
update_ssh_config​(​configfile​, keys)
make_known_hosts_dir​(​ssh_dir​)
# Now let's ensure the services are up and running.
if args.skip_systemd == False:
for env in ENVS:
start_systemd(env)
if __name__ == "__main__":
main()
Powered by Gitiles
txt
json