Coverage for src/usage.py: 100%

1"""To record usage of the app.""" 

2 

3import json 

4import os 

5import tempfile 

6import time 

7from datetime import datetime, timezone 

8from typing import Any 

9 

10import requests 

11 

12from src import get_logger 

13from src.config_parser import get_usage_tracking_enabled, prepare_root_destination 

14 

15LOGGER = get_logger() 

16 

17# Filename for the usage cache. Stored under the root destination 

18# directory (``app.root``) via ``init_cache()``. Kept as a bare 

19# filename so that ``init_cache`` controls where the cache lives. 

20CACHE_FILE_NAME = ".data" 

21NEW_INSTALLATION_ENDPOINT = os.environ.get("NEW_INSTALLATION_ENDPOINT", None) 

22NEW_HEARTBEAT_ENDPOINT = os.environ.get("NEW_HEARTBEAT_ENDPOINT", None) 

23APP_NAME = "icloud-docker" 

24APP_VERSION = os.environ.get("APP_VERSION", "dev") 

25NEW_INSTALLATION_DATA = {"appName": APP_NAME, "appVersion": APP_VERSION} 

26 

27# Retry configuration 

28MAX_RETRIES = int(os.environ.get("USAGE_TRACKING_MAX_RETRIES", "3")) 

29RETRY_BACKOFF_FACTOR = float(os.environ.get("USAGE_TRACKING_RETRY_BACKOFF", "2.0")) 

30 

31 

32def init_cache(config: dict) -> str: 

33 """Initialize the cache file. 

34 

35 Args: 

36 config: Configuration dictionary containing root destination path 

37 

38 Returns: 

39 Absolute path to the cache file 

40 """ 

41 root_destination_path = prepare_root_destination(config=config) 

42 cache_file_path = os.path.join(root_destination_path, CACHE_FILE_NAME) 

43 LOGGER.debug(f"Initialized usage cache at: {cache_file_path}") 

44 return cache_file_path 

45 

46 

47def validate_cache_data(data: dict) -> bool: 

48 """Validate cache data structure. 

49 

50 Args: 

51 data: Dictionary to validate 

52 

53 Returns: 

54 True if data is valid, False otherwise 

55 """ 

56 # Basic structure validation 

57 if not isinstance(data, dict): 

58 return False 

59 

60 # If we have an ID, validate it's a string 

61 if "id" in data and not isinstance(data["id"], str): 

62 return False 

63 

64 # If we have app_version, validate it's a string 

65 if "app_version" in data and not isinstance(data["app_version"], str): 

66 return False 

67 

68 # If we have heartbeat timestamp, validate format. 

69 # Accept both ``%Y-%m-%d %H:%M:%S.%f`` (current) and 

70 # ``%Y-%m-%d %H:%M:%S`` (legacy, microsecond zero) to avoid 

71 # wiping old caches that were written by the str() method. 

72 if "heartbeat_timestamp" in data: 

73 ts = data["heartbeat_timestamp"] 

74 if not isinstance(ts, str): 

75 return False 

76 try: 

77 datetime.strptime(ts, "%Y-%m-%d %H:%M:%S.%f") 

78 except (ValueError, TypeError): 

79 try: 

80 datetime.strptime(ts, "%Y-%m-%d %H:%M:%S") 

81 except (ValueError, TypeError): 

82 return False 

83 

84 return True 

85 

86 

87def load_cache(file_path: str) -> dict: 

88 """Load the cache file with validation and corruption recovery. 

89 

90 Args: 

91 file_path: Absolute path to the cache file 

92 

93 Returns: 

94 Dictionary containing cached usage data 

95 """ 

96 data = {} 

97 if os.path.isfile(file_path): 

98 try: 

99 with open(file_path, encoding="utf-8") as f: 

100 loaded_data = json.load(f) 

101 

102 # Validate the loaded data 

103 if validate_cache_data(loaded_data): 

104 data = loaded_data 

105 LOGGER.debug(f"Loaded and validated usage cache from: {file_path}") 

106 else: 

107 LOGGER.debug(f"Cache data validation failed for {file_path}, starting fresh") 

108 save_cache(file_path=file_path, data={}) 

109 except (json.JSONDecodeError, OSError) as e: 

110 LOGGER.debug(f"Failed to load usage cache from {file_path}: {e}") 

111 LOGGER.debug("Creating new empty cache file due to corruption") 

112 save_cache(file_path=file_path, data={}) 

113 else: 

114 LOGGER.debug(f"Usage cache file not found, creating: {file_path}") 

115 save_cache(file_path=file_path, data={}) 

116 return data 

117 

118 

119def save_cache(file_path: str, data: dict) -> bool: 

120 """Save data to the cache file using atomic operations. 

121 

122 Args: 

123 file_path: Absolute path to the cache file 

124 data: Dictionary containing usage data to save 

125 

126 Returns: 

127 True if save was successful, False otherwise 

128 """ 

129 try: 

130 # Write to temporary file first for atomic operation 

131 dir_name = os.path.dirname(file_path) 

132 with tempfile.NamedTemporaryFile( 

133 mode="w", 

134 encoding="utf-8", 

135 dir=dir_name, 

136 delete=False, 

137 suffix=".tmp", 

138 ) as temp_file: 

139 json.dump(data, temp_file, indent=2) 

140 temp_path = temp_file.name 

141 

142 # Atomically move temp file to final location 

143 os.rename(temp_path, file_path) 

144 LOGGER.debug(f"Atomically saved usage cache to: {file_path}") 

145 return True 

146 except OSError as e: 

147 LOGGER.debug(f"Failed to save usage cache to {file_path}: {e}") 

148 # Clean up temp file if it exists 

149 try: 

150 if "temp_path" in locals(): 

151 os.unlink(temp_path) 

152 except OSError: 

153 pass 

154 return False 

155 

156 

157def post_with_retry( 

158 url: str, 

159 json_data: dict, 

160 timeout: int = 10, 

161 max_retries: int = MAX_RETRIES, 

162 backoff_factor: float = RETRY_BACKOFF_FACTOR, 

163) -> requests.Response | None: 

164 """Post request with exponential backoff retry. 

165 

166 Args: 

167 url: Endpoint URL 

168 json_data: JSON payload 

169 timeout: Request timeout in seconds 

170 max_retries: Maximum number of retry attempts 

171 backoff_factor: Multiplier for exponential backoff 

172 

173 Returns: 

174 Response object if successful, None otherwise 

175 """ 

176 last_exception = None 

177 last_response = None 

178 

179 for attempt in range(max_retries): 

180 try: 

181 response = requests.post(url, json=json_data, timeout=timeout) # type: ignore[arg-type] 

182 

183 # Don't retry on validation errors (4xx except rate limit) 

184 if 400 <= response.status_code < 500 and response.status_code != 429: 

185 LOGGER.debug(f"Non-retriable error (status {response.status_code})") 

186 return response 

187 

188 # Success or retriable error 

189 if response.ok: 

190 return response 

191 

192 # Rate limit (429) or server error (5xx) - retry 

193 last_response = response 

194 LOGGER.debug( 

195 f"Request failed with status {response.status_code}, attempt {attempt + 1}/{max_retries}", 

196 ) 

197 

198 except (requests.ConnectionError, requests.Timeout) as e: 

199 last_exception = e 

200 LOGGER.debug(f"Network error: {e}, attempt {attempt + 1}/{max_retries}") 

201 except Exception as e: 

202 # Catch other exceptions but don't retry 

203 LOGGER.debug(f"Unexpected error during request: {e}") 

204 return None 

205 

206 # Exponential backoff before next retry 

207 if attempt < max_retries - 1: 

208 wait_time = backoff_factor**attempt 

209 LOGGER.debug(f"Waiting {wait_time}s before retry...") 

210 time.sleep(wait_time) 

211 

212 # All retries exhausted — return the last response (if any) so callers 

213 # can distinguish server errors from network failures. 

214 if last_response is not None: 

215 LOGGER.debug(f"All retry attempts failed: HTTP {last_response.status_code}") 

216 return last_response 

217 if last_exception: 

218 LOGGER.debug(f"All retry attempts failed: {last_exception}") 

219 return None 

220 

221 

222def post_new_installation(data: dict, endpoint=NEW_INSTALLATION_ENDPOINT) -> str | None: 

223 """Post new installation to server with retry logic. 

224 

225 Args: 

226 data: Dictionary containing installation data 

227 endpoint: API endpoint URL, defaults to NEW_INSTALLATION_ENDPOINT 

228 

229 Returns: 

230 Installation ID if successful, None otherwise 

231 """ 

232 try: 

233 LOGGER.debug(f"Posting new installation to: {endpoint}") 

234 response = post_with_retry(endpoint, data, timeout=10) 

235 

236 if response and response.ok: 

237 response_data = response.json() 

238 installation_id = response_data["id"] 

239 LOGGER.debug(f"Successfully registered new installation: {installation_id}") 

240 return installation_id 

241 else: 

242 status = response.status_code if response else "no response" 

243 LOGGER.debug(f"Installation registration failed: {status}") 

244 except Exception as e: 

245 LOGGER.debug(f"Failed to post new installation: {e}") 

246 return None 

247 

248 

249def record_new_installation(previous_id: str | None = None) -> str | None: 

250 """Record new or upgrade existing installation. 

251 

252 Args: 

253 previous_id: Previous installation ID for upgrades, None for new installations 

254 

255 Returns: 

256 New installation ID if successful, None otherwise 

257 """ 

258 data = dict(NEW_INSTALLATION_DATA) 

259 if previous_id: 

260 data["previousId"] = previous_id 

261 return post_new_installation(data) 

262 

263 

264def already_installed(cached_data: dict) -> bool: 

265 """Check if already installed. 

266 

267 Args: 

268 cached_data: Dictionary containing cached usage data 

269 

270 Returns: 

271 True if installation is up-to-date, False otherwise 

272 """ 

273 return "id" in cached_data and "app_version" in cached_data and cached_data["app_version"] == APP_VERSION 

274 

275 

276def install(cached_data: dict) -> dict | None: 

277 """Install the app. 

278 

279 Args: 

280 cached_data: Dictionary containing cached usage data 

281 

282 Returns: 

283 Updated cached data dictionary if successful, None otherwise 

284 """ 

285 previous_id = cached_data.get("id", None) 

286 if previous_id: 

287 LOGGER.debug(f"Upgrading existing installation: {previous_id}") 

288 else: 

289 LOGGER.debug("Installing new instance") 

290 

291 new_id = record_new_installation(previous_id) 

292 if new_id: 

293 cached_data["id"] = new_id 

294 cached_data["app_version"] = APP_VERSION 

295 LOGGER.debug(f"Installation completed with ID: {new_id}") 

296 return cached_data 

297 

298 LOGGER.debug("Installation failed") 

299 return None 

300 

301 

302def post_new_heartbeat(data: dict, endpoint=NEW_HEARTBEAT_ENDPOINT) -> bool: 

303 """Post the heartbeat to server with retry logic. 

304 

305 Args: 

306 data: Dictionary containing heartbeat data 

307 endpoint: API endpoint URL, defaults to NEW_HEARTBEAT_ENDPOINT 

308 

309 Returns: 

310 True if heartbeat was sent successfully, False otherwise 

311 """ 

312 try: 

313 LOGGER.debug(f"Posting heartbeat to: {endpoint}") 

314 response = post_with_retry(endpoint, data, timeout=20) 

315 

316 if response and response.ok: 

317 LOGGER.debug("Heartbeat sent successfully") 

318 return True 

319 else: 

320 status = response.status_code if response else "no response" 

321 LOGGER.debug(f"Heartbeat failed: {status}") 

322 except Exception as e: 

323 LOGGER.debug(f"Failed to post heartbeat: {e}") 

324 return False 

325 

326 

327def send_heartbeat(app_id: str | None, data: Any = None) -> bool: 

328 """Prepare and send heartbeat to server. 

329 

330 Args: 

331 app_id: Installation ID for heartbeat identification 

332 data: Additional data to send with heartbeat 

333 

334 Returns: 

335 True if heartbeat was sent successfully, False otherwise 

336 """ 

337 data = {"installationId": app_id, "data": data} 

338 return post_new_heartbeat(data) 

339 

340 

341def _format_timestamp(dt: datetime) -> str: 

342 """Format a datetime as a cache-friendly string. 

343 

344 Uses explicit ``strftime`` so that the microsecond field is always 

345 present (``str(datetime)`` omits it when microseconds are zero, 

346 which would crash ``strptime`` with ``%f`` on load). 

347 """ 

348 return dt.strftime("%Y-%m-%d %H:%M:%S.%f") 

349 

350 

351def _parse_timestamp(ts: str) -> datetime: 

352 """Parse a timestamp string produced by ``_format_timestamp``. 

353 

354 Accepts the ``%f``-fractional format. Returns a timezone-aware UTC 

355 datetime (``tzinfo=timezone.utc``) so it can be compared against 

356 ``current_time()``'s output without naive/aware mismatches. 

357 """ 

358 return datetime.strptime(ts, "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=timezone.utc) 

359 

360 

361def current_time() -> datetime: 

362 """Get current UTC time. 

363 

364 Returns: 

365 Current UTC datetime object (timezone-aware) 

366 """ 

367 return datetime.now(timezone.utc) 

368 

369 

370def heartbeat(cached_data: dict, data: Any) -> dict | None: 

371 """Send heartbeat. 

372 

373 Args: 

374 cached_data: Dictionary containing cached usage data 

375 data: Additional data to send with heartbeat 

376 

377 Returns: 

378 Updated cached data dictionary if heartbeat was sent, 

379 None if heartbeat was throttled or failed 

380 """ 

381 previous_heartbeat = cached_data.get("heartbeat_timestamp", None) 

382 current = current_time() 

383 

384 if previous_heartbeat: 

385 try: 

386 previous = _parse_timestamp(previous_heartbeat) 

387 time_since_last = current - previous 

388 LOGGER.debug(f"Time since last heartbeat: {time_since_last}") 

389 

390 # Check if different UTC day, not just 24 hours 

391 if previous.date() < current.date(): 

392 LOGGER.debug("Sending heartbeat (different UTC day)") 

393 if send_heartbeat(cached_data.get("id"), data=data): 

394 cached_data["heartbeat_timestamp"] = _format_timestamp(current) 

395 return cached_data 

396 else: 

397 LOGGER.debug("Heartbeat send failed") 

398 return None 

399 else: 

400 LOGGER.debug("Heartbeat throttled (same UTC day)") 

401 return None 

402 except ValueError as e: 

403 LOGGER.debug(f"Invalid heartbeat timestamp format: {e}") 

404 # Treat as first heartbeat if timestamp is invalid 

405 

406 # First heartbeat or invalid timestamp 

407 LOGGER.debug("Sending first heartbeat") 

408 if send_heartbeat(cached_data.get("id"), data=data): 

409 cached_data["heartbeat_timestamp"] = _format_timestamp(current) 

410 LOGGER.debug("First heartbeat sent successfully") 

411 return cached_data 

412 else: 

413 LOGGER.debug("First heartbeat send failed") 

414 return None 

415 

416 

417def alive(config: dict | None, data: Any = None) -> bool: 

418 """Record liveliness. 

419 

420 Args: 

421 config: Configuration dictionary (or None to skip tracking) 

422 data: Additional usage data to send with heartbeat 

423 

424 Returns: 

425 True if usage tracking was successful (or skipped), False on failure 

426 """ 

427 # Guard: missing config — skip silently 

428 if config is None: 

429 return True 

430 

431 # Check if usage tracking is disabled 

432 if not get_usage_tracking_enabled(config): 

433 LOGGER.debug("Usage tracking is disabled, skipping") 

434 return True # Return True to not affect main sync loop 

435 

436 LOGGER.debug("Usage tracking alive check started") 

437 

438 cache_file_path = init_cache(config=config) 

439 cached_data = load_cache(cache_file_path) 

440 

441 if not already_installed(cached_data=cached_data): 

442 LOGGER.debug("New installation detected, registering...") 

443 installed_data = install(cached_data=cached_data) 

444 if installed_data is not None: 

445 result = save_cache(file_path=cache_file_path, data=installed_data) 

446 LOGGER.debug("Installation registration completed") 

447 return result 

448 else: 

449 LOGGER.debug("Installation registration failed") 

450 return False 

451 

452 LOGGER.debug("Installation already registered, checking heartbeat") 

453 heartbeat_data = heartbeat(cached_data=cached_data, data=data) 

454 if heartbeat_data is not None: 

455 result = save_cache(file_path=cache_file_path, data=heartbeat_data) 

456 LOGGER.debug("Heartbeat completed successfully") 

457 return result 

458 

459 LOGGER.debug("No heartbeat required or heartbeat failed") 

460 return True