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))