@cryptotaxi247 / infra-1 / commits / 1282e2ec

hydra-packet-importer: format with black, check with mypy at build time

Graham Christensen committed Jun 3, 2022 at 10:19 UTC 1282e2ec69485b26fdc83b9f31e43826e5cd3d9c
2 files changed +95 -63
hydra-packet-importer/default.nix
+7
@@ -5,11 +5,18 @@ python3.pkgs.buildPythonApplication {
5
6 format = "other";
7
8 + nativeBuildInputs = [
9 + python3.pkgs.mypy
10 + python3.pkgs.black
11 + ];
12 +
13 propagatedBuildInputs = [
14 python3.pkgs.packet-python
15 ];
16
17 installPhase = ''
18 + mypy ./import.py
19 + black --check ./import.py
20 mkdir -p $out/bin
21 mv import.py $out/bin/hydra-packet-importer
22 '';
hydra-packet-importer/import.py
+88 -63
@@ -1,7 +1,7 @@
1 #!/usr/bin/env python3
2
3 import json
4 -import packet
4 +import packet # type: ignore
5 import base64
6 from pprint import pprint
7 import subprocess
@@ -12,12 +12,12 @@ DeviceKeys = List[Dict[str, Any]]
12
13 if sys.version_info >= (3, 8):
14 from typing import TypedDict
15 +
16 class Metadata(TypedDict):
17 features: List[str]
18 max_jobs: int
19 system_types: List[str]
20
20 -
21 class RemoteBuilder(TypedDict):
22 metadata: Metadata
23 ssh_key: str
@@ -32,12 +32,15 @@ if sys.version_info >= (3, 8):
32 system: str
33 port: int
34 key: str
35 +
36 +
37 else:
38 Device = Dict[str, Any]
39 RemoteBuilder = Dict[str, Any]
40 HostKey = Dict[str, Any]
41 Metadata = Dict[str, Any]
42
43 +
44 def debug(*args: Any, **kwargs: Any) -> None:
45 print(*args, file=sys.stderr, **kwargs)
46
@@ -45,42 +48,45 @@ def debug(*args: Any, **kwargs: Any) -> None:
48 def get_devices(manager: Any) -> List[Device]:
49 devices: List[Device] = []
50
48 - page: Optional[str] = 'projects/{}/devices?page={}'.format(config['project_id'], 1)
51 + page: Optional[str] = "projects/{}/devices?page={}".format(config["project_id"], 1)
52 while page is not None:
53 debug(page)
54 data: Dict[str, Any] = manager.call_api(page)
52 - if data['meta']['next'] is None:
55 + if data["meta"]["next"] is None:
56 page = None
57 else:
55 - page = data['meta']['next']['href']
58 + page = data["meta"]["next"]["href"]
59
57 - for device in data['devices']:
58 - if device['state'] != 'active':
60 + for device in data["devices"]:
61 + if device["state"] != "active":
62 continue
63
61 - if not set(config['mandatory_tags']).issubset(device['tags']):
64 + if not set(config["mandatory_tags"]).issubset(device["tags"]):
65 continue
66
64 - if not set(device['tags']).isdisjoint(config['skip_tags']):
67 + if not set(device["tags"]).isdisjoint(config["skip_tags"]):
68 continue
69
67 - remote_builder_info = get_remote_builder_info(manager, device['id'])
70 + remote_builder_info = get_remote_builder_info(manager, device["id"])
71 if remote_builder_info is None:
72 continue
73
71 - devices.append({
72 - "hostname": device['hostname'],
73 - "address": "{}.packethost.net".format(device['short_id']),
74 - "type": device['plan']['name'],
75 - "remote_builder_info": remote_builder_info,
76 - })
74 + devices.append(
75 + {
76 + "hostname": device["hostname"],
77 + "address": "{}.packethost.net".format(device["short_id"]),
78 + "type": device["plan"]["name"],
79 + "remote_builder_info": remote_builder_info,
80 + }
81 + )
82
83 return devices
84
85 +
86 def get_remote_builder_info(manager, device_id: str) -> Union[RemoteBuilder, str, None]:
87 # ... 50 is probably enough.
88 try:
83 - events_url = 'devices/{}/events?per_page=50'.format(device_id)
89 + events_url = "devices/{}/events?per_page=50".format(device_id)
90 debug(events_url)
91 data = manager.call_api(events_url)
92 except:
@@ -90,8 +96,8 @@ def get_remote_builder_info(manager, device_id: str) -> Union[RemoteBuilder, str
96 host_key: Optional[HostKey] = None
97 ssh_key: Optional[str] = None
98 metadata: Optional[Metadata] = None
93 - for event in data['events']:
94 - if event['type'] == 'provisioning.104.01':
99 + for event in data["events"]:
100 + if event["type"] == "provisioning.104.01":
101 # we reached a "Device connected to DHCP system" event,
102 # indicating a reboot.
103 #
@@ -102,25 +108,28 @@ def get_remote_builder_info(manager, device_id: str) -> Union[RemoteBuilder, str
108 # If we receive a LOT of spam (> 50 spams!) like that, we
109 # will return None because we never reach this message.
110 if host_key is not None:
105 - key = strip_ssh_key_comment(host_key['key'])
111 + key = strip_ssh_key_comment(host_key["key"])
112 if key is not None:
113 if metadata is not None:
108 - return { "metadata": metadata, "ssh_key": key }
114 + return {"metadata": metadata, "ssh_key": key}
115 else:
116 return key
117 else:
118 return ssh_key
113 - if event['type'] == 'user.1001':
119 + if event["type"] == "user.1001":
120 try:
115 - host_keys: List[HostKey] = [key for key in json.loads(event['body']) if key['port'] == 22]
121 + host_keys: List[HostKey] = [
122 + key for key in json.loads(event["body"]) if key["port"] == 22
123 + ]
124 host_key = host_keys[0]
125 except:
118 - ssh_key = strip_ssh_key_comment(event['body'])
119 - if event['type'] == 'user.1002':
120 - metadata = json.loads(event['body'])
126 + ssh_key = strip_ssh_key_comment(event["body"])
127 + if event["type"] == "user.1002":
128 + metadata = json.loads(event["body"])
129
130 return None
131
132 +
133 def strip_ssh_key_comment(key: str) -> Optional[str]:
134 ssh_key_parts = key.rsplit(" ", 1)
135 if len(ssh_key_parts) == 2:
@@ -132,62 +141,78 @@ def strip_ssh_key_comment(key: str) -> Optional[str]:
141
142 def main(config: Dict[str, Any]) -> None:
143 rows = []
135 - manager = packet.Manager(auth_token=config['token'])
144 + manager = packet.Manager(auth_token=config["token"])
145 found = 0
146 for device in get_devices(manager):
147 found += 1
139 - debug("# {} ({})".format(device['hostname'], device['address']))
140 - if device['type'] not in config['plans']:
141 - debug("# Skipping {} (type {}) as it has no configured plan".format(
142 - device['hostname'],
143 - device['type'])
148 + debug("# {} ({})".format(device["hostname"], device["address"]))
149 + if device["type"] not in config["plans"]:
150 + debug(
151 + "# Skipping {} (type {}) as it has no configured plan".format(
152 + device["hostname"], device["type"]
153 + )
154 )
155 continue
156
147 - builder_info = device['remote_builder_info']
148 - default_stats = config['plans'][device['type']]
149 - if device['hostname'] in config['name_overrides']:
150 - specific_stats = config['name_overrides'][device['hostname']]
157 + builder_info = device["remote_builder_info"]
158 + default_stats = config["plans"][device["type"]]
159 + if device["hostname"] in config["name_overrides"]:
160 + specific_stats = config["name_overrides"][device["hostname"]]
161 else:
162 specific_stats = {}
153 - lookup = lambda key: specific_stats.get(key, device.get(key, default_stats.get(key)))
163 + lookup = lambda key: specific_stats.get(
164 + key, device.get(key, default_stats.get(key))
165 + )
166
155 - lookup_default = lambda key, default: default if not lookup(key) else lookup(key)
167 + lookup_default = (
168 + lambda key, default: default if not lookup(key) else lookup(key)
169 + )
170
171 if isinstance(builder_info, str):
172 key = builder_info
173 # root@address system,list /var/lib/ssh.key maxJobs speedFactor feature,list mandatory,features public-host-key
160 - rows.append(" ".join([
161 - "{user}@{host}".format(user=lookup("user"),host=lookup("address")),
162 - ",".join(lookup("system_types")),
163 - str(lookup("ssh_key")),
164 - str(lookup("max_jobs")),
165 - str(lookup("speed_factor")),
166 - ",".join(lookup_default("features", ["-"])),
167 - ",".join(lookup_default("mandatory_features", ["-"])),
168 - base64.b64encode(key.encode()).decode("utf-8")
169 - ]))
174 + rows.append(
175 + " ".join(
176 + [
177 + "{user}@{host}".format(
178 + user=lookup("user"), host=lookup("address")
179 + ),
180 + ",".join(lookup("system_types")),
181 + str(lookup("ssh_key")),
182 + str(lookup("max_jobs")),
183 + str(lookup("speed_factor")),
184 + ",".join(lookup_default("features", ["-"])),
185 + ",".join(lookup_default("mandatory_features", ["-"])),
186 + base64.b64encode(key.encode()).decode("utf-8"),
187 + ]
188 + )
189 + )
190 else:
191 # root@address system,list /var/lib/ssh.key maxJobs speedFactor feature,list mandatory,features public-host-key
172 - rows.append(" ".join([
173 - "{user}@{host}".format(user=lookup("user"),host=lookup("address")),
174 - ",".join(builder_info['metadata']['system_types']),
175 - str(lookup("ssh_key")),
176 - str(builder_info['metadata']['max_jobs']),
177 - str(lookup("speed_factor")),
178 - ",".join(builder_info['metadata']['features']),
179 - ",".join(lookup_default("mandatory_features", ["-"])),
180 - base64.b64encode(builder_info['ssh_key'].encode()).decode("utf-8")
181 - ]))
182 -
183 -
184 -
192 + rows.append(
193 + " ".join(
194 + [
195 + "{user}@{host}".format(
196 + user=lookup("user"), host=lookup("address")
197 + ),
198 + ",".join(builder_info["metadata"]["system_types"]),
199 + str(lookup("ssh_key")),
200 + str(builder_info["metadata"]["max_jobs"]),
201 + str(lookup("speed_factor")),
202 + ",".join(builder_info["metadata"]["features"]),
203 + ",".join(lookup_default("mandatory_features", ["-"])),
204 + base64.b64encode(builder_info["ssh_key"].encode()).decode(
205 + "utf-8"
206 + ),
207 + ]
208 + )
209 + )
210
186 - debug("# {} / {}".format(len(rows),found))
211 + debug("# {} / {}".format(len(rows), found))
212 print("\n".join(rows))
213
214 +
215 if __name__ == "__main__":
216 with open(sys.argv[1]) as config_file:
217 config = json.load(config_file)
218 main(config)
193 -