@cryptotaxi247 / infra / commits / 431df21c

hydra-packet-importer: let machines describe themselves

Graham Christensen committed Mar 7, 2020 at 22:21 UTC 431df21c24e9da9206896e8d45fe939ac99b6b8b
1 file changed +89 -38
hydra-packet-importer/import.py
+89 -38
@@ -8,8 +8,35 @@ import subprocess
8 import sys
9 from typing import Union, Dict, Any, List, Optional
10
11 -Device = Dict[str, Any]
12 -DeviceKeys =Union[None, List[Dict[str, Any]], str]
11 +DeviceKeys = List[Dict[str, Any]]
12 +
13 +if sys.version_info >= (3, 8):
14 + from typing import TypedDict
15 + class Metadata(TypedDict):
16 + features: List[str]
17 + max_jobs: int
18 + system_types: List[str]
19 +
20 +
21 + class RemoteBuilder(TypedDict):
22 + metadata: Metadata
23 + ssh_key: str
24 +
25 + class Device(TypedDict):
26 + hostname: str
27 + address: str
28 + type: str
29 + remote_builder_info: Union[RemoteBuilder, str]
30 +
31 + class HostKey(TypedDict):
32 + system: str
33 + port: int
34 + key: str
35 +else:
36 + Device = Dict[str, Any]
37 + RemoteBuilder = Dict[str, Any]
38 + HostKey = Dict[str, Any]
39 + Metadata = Dict[str, Any]
40
41 def debug(*args: Any, **kwargs: Any) -> None:
42 print(*args, file=sys.stderr, **kwargs)
@@ -37,26 +64,32 @@ def get_devices(manager: Any) -> List[Device]:
64 if not set(device['tags']).isdisjoint(config['skip_tags']):
65 continue
66
40 - host_key = get_device_key(manager, device)
41 - if host_key is None:
67 + remote_builder_info = get_remote_builder_info(manager, device['id'])
68 + if remote_builder_info is None:
69 continue
70
71 devices.append({
72 "hostname": device['hostname'],
73 "address": "{}.packethost.net".format(device['short_id']),
74 "type": device['plan']['name'],
48 - "host_key": host_key,
75 + "remote_builder_info": remote_builder_info,
76 })
77
78 return devices
79
53 -def get_device_key(manager, device: Device) -> DeviceKeys:
80 +def get_remote_builder_info(manager, device_id: str) -> Union[RemoteBuilder, str, None]:
81 # ... 50 is probably enough.
55 - events_url = 'devices/{}/events?per_page=50'.format(device['id'])
56 - debug(events_url)
57 - data = manager.call_api(events_url)
58 -
59 - ssh_key = None
82 + try:
83 + events_url = 'devices/{}/events?per_page=50'.format(device_id)
84 + debug(events_url)
85 + data = manager.call_api(events_url)
86 + except:
87 + # 404 probably
88 + return None
89 +
90 + host_key: Optional[HostKey] = None
91 + ssh_key: Optional[str] = None
92 + metadata: Optional[Metadata] = None
93 for event in data['events']:
94 if event['type'] == 'provisioning.104.01':
95 # we reached a "Device connected to DHCP system" event,
@@ -68,19 +101,35 @@ def get_device_key(manager, device: Device) -> DeviceKeys:
101 #
102 # If we receive a LOT of spam (> 50 spams!) like that, we
103 # will return None because we never reach this message.
71 - return ssh_key
104 + if host_key is not None:
105 + key = strip_ssh_key_comment(host_key['key'])
106 + if key is not None:
107 + if metadata is not None:
108 + return { "metadata": metadata, "ssh_key": key }
109 + else:
110 + return key
111 + else:
112 + return ssh_key
113 if event['type'] == 'user.1001':
114 try:
74 - ssh_key = json.loads(event['body'])
115 + host_keys: List[HostKey] = [key for key in json.loads(event['body']) if key['port'] == 22]
116 + host_key = host_keys[0]
117 except:
76 - ssh_key_parts = event['body'].rsplit(" ", 1)
77 - if len(ssh_key_parts) == 2:
78 - ssh_key = ssh_key_parts[0] + "\n"
79 - else:
80 - debug("# Skipped due keyscan failed to split on ' '")
118 + ssh_key = strip_ssh_key_comment(event['body'])
119 + if event['type'] == 'user.1002':
120 + metadata = json.loads(event['body'])
121
122 return None
123
124 +def strip_ssh_key_comment(key: str) -> Optional[str]:
125 + ssh_key_parts = key.rsplit(" ", 1)
126 + if len(ssh_key_parts) == 2:
127 + return ssh_key_parts[0]
128 + else:
129 + debug("# Skipped due keyscan failed to split on ' '")
130 + return None
131 +
132 +
133 def main(config: Dict[str, Any]) -> None:
134 rows = []
135 manager = packet.Manager(auth_token=config['token'])
@@ -95,33 +144,20 @@ def main(config: Dict[str, Any]) -> None:
144 )
145 continue
146
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']]
151 else:
152 specific_stats = {}
103 -
153 lookup = lambda key: specific_stats.get(key, device.get(key, default_stats.get(key)))
105 - lookup_default = lambda key, default: default if not lookup(key) else lookup(key)
106 -
107 - keys: DeviceKeys = device["host_key"]
108 - key: Optional[str] = None
109 - if keys is None:
110 - debug("# no key data")
111 - elif isinstance(keys, str):
112 - key = keys
113 - else:
114 - keys_matching = [keyrec["key"] for keyrec in keys
115 - if keyrec['system'] in lookup("system_types")]
116 - keys_matching = sorted(keys_matching)
117 - if len(keys_matching) > 0:
118 - key = keys_matching[0]
154
120 - if key is None:
121 - debug("# no matching data")
155 + lookup_default = lambda key, default: default if not lookup(key) else lookup(key)
156
123 - # root@address system,list /var/lib/ssh.key maxJobs speedFactor feature,list mandatory,features public-host-key
124 - rows.append(" ".join([
157 + if isinstance(builder_info, str):
158 + key = builder_info
159 + # 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")),
@@ -130,7 +166,22 @@ def main(config: Dict[str, Any]) -> None:
166 ",".join(lookup_default("features", ["-"])),
167 ",".join(lookup_default("mandatory_features", ["-"])),
168 base64.b64encode(key.encode()).decode("utf-8")
133 - ]))
169 + ]))
170 + else:
171 + # 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 +
185
186 debug("# {} / {}".format(len(rows),found))
187 print("\n".join(rows))