Coverage for src/lib/address_book.py: 82%

1 

2import datetime as dt 

3from os import path, remove 

4from logging import getLogger, Logger 

5from sty import fg 

6from lib.client import Client 

7 

8from lib.overlay import Node, Distance 

9from lib.helper import read_json_file, write_json_file 

10from lib.contact import Contact 

11 

12 

13class AddressBook(): 

14 _path: str 

15 _config: dict 

16 _ab_config: dict 

17 _clients_by_uuid: dict[str, Client] 

18 _clients_by_id: dict[str, Client] 

19 _changes: bool 

20 _clients_ttl: dt.timedelta 

21 _logger: Logger 

22 

23 def __init__(self, path: str, config: dict = None): 

24 self._path = path 

25 self._config = config 

26 self._ab_config = self._config['address_book'] 

27 self._clients_by_uuid = dict() 

28 self._clients_by_id = dict() 

29 self._changes = False 

30 

31 self._logger = getLogger('app.address_book') 

32 self._logger.info('init()') 

33 

34 if self._ab_config is None: 

35 self._clients_ttl = dt.timedelta(hours=1) 

36 else: 

37 self._clients_ttl = dt.timedelta(hours=self._ab_config['client_retention_time']) 

38 

39 self._logger.info('clients_ttl %s', self._clients_ttl) 

40 

41 def load(self): 

42 self._logger.debug('load()') 

43 

44 _data = read_json_file(self._path, {}) 

45 for client_uuid, row in _data.items(): 

46 client = Client() 

47 client.uuid = client_uuid 

48 client.from_dict(row) 

49 

50 self._logger.debug('load client: %s', client) 

51 

52 self._clients_by_uuid[client_uuid] = client 

53 if client.id is not None: 

54 if client.id in self._clients_by_id: 

55 self._logger.warning('Client ID already exists: %s', client.id) 

56 

57 self._clients_by_id[client.id] = client 

58 

59 key_file_path = path.join(self._config['keys_dir'], client.id + '.pem') 

60 if path.isfile(key_file_path): 

61 client.load_public_key_from_pem_file(key_file_path) 

62 

63 def save(self) -> bool: 

64 self._logger.debug('save() changes=%s', self._changes) 

65 

66 if not self._changes: 

67 return False 

68 

69 _data = dict() 

70 for client_uuid, client in self._clients_by_uuid.items(): 

71 self._logger.debug('save client: %s', client) 

72 

73 _data[client_uuid] = client.as_dict() 

74 

75 if client.id is not None: 

76 key_file_path = path.join(self._config['keys_dir'], client.id + '.pem') 

77 if not path.isfile(key_file_path): 

78 client.write_public_key_to_pem_file(key_file_path) 

79 

80 write_json_file(self._path, _data) 

81 self._changes = False 

82 

83 return True 

84 

85 def get_clients(self) -> dict[str, Client]: 

86 return self._clients_by_uuid 

87 

88 def get_clients_len(self) -> int: 

89 return len(self._clients_by_uuid) 

90 

91 def get_bootstrap_clients(self) -> list: 

92 def ffunc(_client): 

93 return _client[1].is_bootstrap 

94 bootstrap_clients = list(filter(ffunc, self._clients_by_uuid.items())) 

95 return bootstrap_clients 

96 

97 def get_bootstrap_clients_len(self) -> int: 

98 def ffunc(_client): 

99 return _client[1].is_bootstrap 

100 bootstrap_clients = list(filter(ffunc, self._clients_by_uuid.items())) 

101 return len(bootstrap_clients) 

102 

103 def get_client_by_id(self, id: str) -> Client: 

104 self._logger.debug('get_client_by_id(%s)', id) 

105 

106 if id in self._clients_by_id: 

107 return self._clients_by_id[id] 

108 

109 return None 

110 

111 def get_client_by_uuid(self, uuid: str) -> Client: 

112 self._logger.debug('get_client_by_uuid(%s)', uuid) 

113 

114 if uuid in self._clients_by_uuid: 

115 return self._clients_by_uuid[uuid] 

116 

117 return None 

118 

119 def get_client_by_addr_port(self, addr: str, port: int): 

120 def ffunc(_client): 

121 return _client[1].address == addr and _client[1].port == port 

122 _clients = list(filter(ffunc, self._clients_by_uuid.items())) 

123 

124 if len(_clients) > 0: 

125 return _clients[0][1] 

126 

127 return None 

128 

129 def add_client(self, id: str = None, addr: str = None, port: int = None) -> Client: 

130 self._logger.debug('add_client(%s, %s, %s)', id, addr, port) 

131 

132 client = Client() 

133 

134 if id is not None: 

135 client.set_id(id) 

136 

137 if addr is not None: 

138 client.address = addr 

139 

140 if port is not None: 

141 client.port = port 

142 

143 self._clients_by_uuid[client.uuid] = client 

144 if client.id is not None: 

145 self._clients_by_id[client.id] = client 

146 

147 self.changed() 

148 

149 return client 

150 

151 def append_client(self, client: Client): 

152 self._logger.debug('append_client(%s)', client) 

153 

154 self._clients_by_uuid[client.uuid] = client 

155 if client.id is not None: 

156 self._clients_by_id[client.id] = client 

157 

158 self.changed() 

159 

160 def remove_client(self, client: Client, force: bool = False) -> bool: 

161 self._logger.debug('remove_client(%s, %s)', client, force) 

162 

163 if not force and client.is_trusted: 

164 return False 

165 

166 key_file_path = path.join(self._config['keys_dir'], client.id + '.pem') 

167 if path.isfile(key_file_path): 

168 remove(key_file_path) 

169 

170 del self._clients_by_uuid[client.uuid] 

171 if client.id is not None: 

172 del self._clients_by_id[client.id] 

173 

174 self.changed() 

175 

176 return True 

177 

178 def add_bootstrap(self, file_path: str): 

179 self._logger.debug('add_bootstrap(%s)', file_path) 

180 _data = read_json_file(file_path, []) 

181 for row in _data: 

182 contact = Contact.parse(row) 

183 

184 _client = self.get_client_by_addr_port(contact.addr, contact.port) 

185 if _client is not None: 

186 self._logger.debug('bootstrap client already exists: %s', _client) 

187 continue 

188 

189 client = Client() 

190 client.address = contact.addr 

191 client.port = contact.port 

192 client.is_bootstrap = True 

193 client.debug_add = 'bootstrap' 

194 

195 self._clients_by_uuid[client.uuid] = client 

196 if client.id is not None: 

197 self._clients_by_id[client.id] = client 

198 

199 self.changed() 

200 

201 write_json_file(file_path, []) 

202 

203 def changed(self): 

204 self._changes = True 

205 

206 def hard_clean_up(self, local_id: str = None): 

207 self._logger.debug('hard_clean_up(%s)', local_id) 

208 

209 # remove local_id 

210 if local_id is not None and local_id in self._clients_by_id: 

211 del self._clients_by_id[local_id] 

212 

213 _clients = list(self._clients_by_uuid.values()) 

214 _clients_len = len(_clients) 

215 

216 if _clients_len <= self._ab_config['max_clients']: 

217 return 

218 

219 # remove bootstrap clients with no meetings 

220 _clients = list(filter(lambda _client: _client.is_bootstrap and _client.meetings == 0, _clients)) 

221 for client in _clients: 

222 if self.remove_client(client): 

223 _clients_len -= 1 

224 if _clients_len <= self._ab_config['max_clients']: 

225 return 

226 

227 # remove out-of-date clients (invalid client_retention_time) 

228 _clients = list(self._clients_by_uuid.values()) 

229 _clients = list(filter(lambda _client: dt.datetime.now(dt.UTC) - _client.used_at > self._clients_ttl, _clients)) 

230 _clients.sort(key=lambda _client: _client.used_at) 

231 for client in _clients: 

232 if self.remove_client(client): 

233 _clients_len -= 1 

234 if _clients_len <= self._ab_config['max_clients']: 

235 return 

236 

237 # remove clients, sorted by meetings 

238 _clients = list(self._clients_by_uuid.values()) 

239 _clients.sort(key=lambda _client: _client.meetings) 

240 

241 for client in _clients: 

242 if self.remove_client(client): 

243 _clients_len -= 1 

244 if _clients_len <= self._ab_config['max_clients']: 

245 return 

246 

247 def soft_clean_up(self, local_id: str = None): 

248 self._logger.debug('soft_clean_up(%s)', local_id) 

249 

250 # remove local_id 

251 if local_id is not None and local_id in self._clients_by_id: 

252 del self._clients_by_id[local_id] 

253 

254 _clients = list(self._clients_by_uuid.values()) 

255 _clients_len = len(_clients) 

256 

257 # remove out-of-date clients (invalid client_retention_time) 

258 _clients = list(self._clients_by_uuid.values()) 

259 _clients = list(filter(lambda _client: dt.datetime.now(dt.UTC) - _client.used_at > self._clients_ttl and _client.meetings == 0, _clients)) 

260 _clients.sort(key=lambda _client: _client.used_at) 

261 for client in _clients: 

262 if self.remove_client(client): 

263 _clients_len -= 1 

264 if _clients_len <= self._ab_config['max_clients']: 

265 return 

266 

267 def get_nearest_to(self, node: Node, limit: int = 20, with_contact_infos: bool = None) -> list: 

268 def sort_key(_client: Client) -> Distance: 

269 print(f'-> sort_key client: {_client}') 

270 print(f'-> sort_key node: {_client.node}') 

271 return _client.node.distance(node) 

272 

273 _clients = list(self._clients_by_uuid.values()) 

274 _clients.sort(key=sort_key) 

275 

276 if with_contact_infos: 

277 _clients = list(filter(lambda _client: with_contact_infos == _client.has_contact(), _clients)) 

278 

279 return _clients[:limit]