gerrit.wikimedia.org
Gitiles
Code ReviewSign In
gerrit.wikimedia.org / operations / debs / wmf-sre-laptop / ab18ff7a42963c6d9e83ff6e6bdca3a286651815 / . / scripts / wmf-update-ssh-config
blob: 0212b660a629e9be94237181a3d2e0e9ad053fe9 [file] [log] [blame]
Giuseppe Lavagetto612c9422020-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 Caro406635d2021-03-02 17:33:28 +0100[diff] [blame]
11
CONF_MARKER = "### wmf-laptop-sre - do not edit below this point ###\n"
Giuseppe Lavagetto612c9422020-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 Caro406635d2021-03-02 17:33:28 +0100[diff] [blame]
29
with open​(​"/usr/share/wmf-laptop-sre/ssh-client-config"​, "r") as fh:
Giuseppe Lavagetto612c9422020-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()
Powered by Gitiles
txt
json