packet importer: handle multi-key format
Eelco Dolstra committed
Mar 5, 2020 at 05:42 UTC
3ac98e3264f33c4ba31ca8dc3c1d20359eb3f955
1 file changed
+37
-13
hydra-packet-importer/import.py
+37
-13
@@ -6,18 +6,22 @@ import base64
6
from pprint import pprint
7
import subprocess
8
import sys
9
+from typing import Union, Dict, Any, List, Optional
10
10
-def debug(*args, **kwargs):
11
+Device = Dict[str, Any]
12
+DeviceKeys =Union[None, List[Dict[str, Any]], str]
13
+
14
+def debug(*args: Any, **kwargs: Any) -> None:
15
print(*args, file=sys.stderr, **kwargs)
16
17
14
-def get_devices(manager):
15
- devices = []
18
+def get_devices(manager: Any) -> List[Device]:
19
+ devices: List[Device] = []
20
17
- page = 'projects/{}/devices?page={}'.format(config['project_id'], 1)
21
+ page: Optional[str] = 'projects/{}/devices?page={}'.format(config['project_id'], 1)
22
while page is not None:
23
debug(page)
20
- data = manager.call_api(page)
24
+ data: Dict[str, Any] = manager.call_api(page)
25
if data['meta']['next'] is None:
26
page = None
27
else:
@@ -46,7 +50,7 @@ def get_devices(manager):
50
51
return devices
52
49
-def get_device_key(manager, device):
53
+def get_device_key(manager, device: Device) -> DeviceKeys:
54
# ... 50 is probably enough.
55
events_url = 'devices/{}/events?per_page=50'.format(device['id'])
56
debug(events_url)
@@ -66,15 +70,18 @@ def get_device_key(manager, device):
70
# will return None because we never reach this message.
71
return ssh_key
72
if event['type'] == 'user.1001':
69
- ssh_key_parts = event['body'].rsplit(" ", 1)
70
- if len(ssh_key_parts) == 2:
71
- ssh_key = ssh_key_parts[0] + "\n"
72
- else:
73
- debug("# Skipped due keyscan failed to split on ' '")
73
+ try:
74
+ ssh_key = json.loads(event['body'])
75
+ 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 ' '")
81
82
return None
83
77
-def main(config):
84
+def main(config: Dict[str, Any]) -> None:
85
rows = []
86
manager = packet.Manager(auth_token=config['token'])
87
found = 0
@@ -97,6 +104,22 @@ def main(config):
104
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]
119
+
120
+ if key is None:
121
+ debug("# no matching data")
122
+
123
# root@address system,list /var/lib/ssh.key maxJobs speedFactor feature,list mandatory,features public-host-key
124
rows.append(" ".join([
125
"{user}@{host}".format(user=lookup("user"),host=lookup("address")),
@@ -106,7 +129,7 @@ def main(config):
129
str(lookup("speed_factor")),
130
",".join(lookup_default("features", ["-"])),
131
",".join(lookup_default("mandatory_features", ["-"])),
109
- base64.b64encode(device['host_key'].encode()).decode("utf-8")
132
+ base64.b64encode(key.encode()).decode("utf-8")
133
]))
134
135
debug("# {} / {}".format(len(rows),found))
@@ -116,3 +139,4 @@ if __name__ == "__main__":
139
with open(sys.argv[1]) as config_file:
140
config = json.load(config_file)
141
main(config)
142
+