Coverage for src/twofas/cli.py: 26%

166 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-01-29 11:11 +0100

1""" 

2This file contains the Typer CLI. 

3""" 

4# pragma: no cover 

5 

6import os 

7import sys 

8import typing 

9 

10import questionary 

11import rich 

12import typer 

13from lib2fas._security import keyring_manager 

14from lib2fas._types import TwoFactorAuthDetails 

15from lib2fas.core import TwoFactorStorage, load_services 

16 

17from .__about__ import __version__ 

18from .cli_settings import ( 

19 expand_path, 

20 get_cli_setting, 

21 load_cli_settings, 

22 set_cli_setting, 

23) 

24from .cli_support import ( 

25 clear, 

26 exit_with_clear, 

27 generate_choices, 

28 generate_custom_style, 

29 state, 

30) 

31 

32app = typer.Typer() 

33 

34TwoFactorDetailStorage: typing.TypeAlias = TwoFactorStorage[TwoFactorAuthDetails] 

35 

36 

37def prepare_to_generate(filename: str = None) -> TwoFactorDetailStorage: 

38 """ 

39 Clear old keyring entries (from previous sessions) and decrypt the selected 2fas file. 

40 """ 

41 keyring_manager.cleanup_keyring() 

42 return load_services(filename or default_2fas_file()) 

43 

44 

45def print_for_service(service: TwoFactorAuthDetails) -> None: 

46 """ 

47 Print the name, current TOTP code and optionally username for a specific service. 

48 """ 

49 service_name = service.name 

50 code = service.generate() 

51 

52 if state.verbose: 

53 username = service.otp.account # or .label ? 

54 rich.print(f"- {service_name} ({username}): {code}") 

55 else: 

56 rich.print(f"- {service_name}: {code}") 

57 

58 

59def generate_all_totp(services: TwoFactorDetailStorage) -> None: 

60 """ 

61 Generate TOTP codes for all services. 

62 """ 

63 for service in services: 

64 print_for_service(service) 

65 

66 

67def generate_one_otp(services: TwoFactorDetailStorage) -> None: 

68 """ 

69 Query the user for a service, then generate a TOTP code for it. 

70 """ 

71 service_name: str 

72 while service_name := questionary.autocomplete( 

73 "Choose a service", choices=services.keys(), style=generate_custom_style() 

74 ).ask(): 

75 for service in services.find(service_name): 

76 print_for_service(service) 

77 

78 

79@clear 

80def show_service_info(services: TwoFactorDetailStorage, about: str) -> None: 

81 """ 

82 `--info <service>` to show the raw JSON info for a service as stored in the .2fas file. 

83 """ 

84 rich.print(services[about]) 

85 

86 

87def show_service_info_interactive(services: TwoFactorDetailStorage) -> None: 

88 """ 

89 Menu when choosing "Info about a Service". 

90 

91 The raw JSON info for a service as stored in the .2fas file will be printed out. 

92 """ 

93 about: str 

94 while about := questionary.select( 

95 "About which service?", choices=services.keys(), style=generate_custom_style() 

96 ).ask(): 

97 show_service_info(services, about) 

98 if questionary.press_any_key_to_continue("Press 'Enter' to continue; Other keys to exit").ask() is None: 

99 exit_with_clear(0) 

100 

101 

102@clear 

103def command_interactive(filename: str = None) -> None: 

104 """ 

105 Interactive menu when using 2fas without any action flags. 

106 """ 

107 if not filename: 

108 # get from settings or 

109 filename = default_2fas_file() 

110 

111 services = prepare_to_generate(filename) 

112 

113 rich.print(f"Active file: [blue]{filename}[/blue]") 

114 

115 match questionary.select( 

116 "What do you want to do?", 

117 choices=generate_choices( 

118 { 

119 "Generate a TOTP code": "generate-one", 

120 "Generate all TOTP codes": "generate-all", 

121 "Info about a Service": "see-info", 

122 "Settings": "settings", 

123 } 

124 ), 

125 use_shortcuts=True, 

126 style=generate_custom_style(), 

127 ).ask(): 

128 case "generate-one": 

129 # query list of items 

130 return generate_one_otp(services) 

131 case "generate-all": 

132 # show all 

133 return generate_all_totp(services) 

134 case "see-info": 

135 return show_service_info_interactive(services) 

136 case "settings": 

137 return command_settings(filename) 

138 case _: 

139 exit_with_clear(0) 

140 

141 

142def add_2fas_file() -> str: 

143 """ 

144 Query the user for a 2fas file and remember it for later. 

145 """ 

146 settings = state.settings 

147 

148 filename: str = questionary.path( 

149 "Path to .2fas file?", 

150 validate=lambda it: it.endswith(".2fas"), 

151 # file_filter=lambda it: it.endswith(".2fas"), 

152 style=generate_custom_style(), 

153 ).ask() 

154 

155 if filename is None: 

156 return exit_with_clear(0) 

157 

158 filename = expand_path(filename) 

159 

160 settings.add_file(filename) 

161 return filename 

162 

163 

164def default_2fas_file() -> str: 

165 """ 

166 Load the default 2fas file from settings or query the user for it. 

167 """ 

168 settings = state.settings 

169 if settings.default_file: 

170 return settings.default_file 

171 

172 elif settings.files: 

173 return settings.files[0] 

174 

175 filename = add_2fas_file() 

176 set_cli_setting("default-file", filename) 

177 

178 return expand_path(filename) 

179 

180 

181def default_2fas_services() -> TwoFactorDetailStorage: 

182 """ 

183 Load the 2fas services from the active default file. 

184 """ 

185 filename = default_2fas_file() 

186 return prepare_to_generate(filename) 

187 

188 

189def command_generate(filename: str | None, other_args: list[str]) -> None: 

190 """ 

191 Handles the generation of OTP codes for the specified service(s) \ 

192 or initiates an interactive menu if no services are specified. 

193 

194 Args: 

195 filename: path to the active .2fas file 

196 other_args: list of services to generate codes for. If empty, an interactive menu will be shown. 

197 """ 

198 storage = prepare_to_generate(filename) 

199 found: list[TwoFactorAuthDetails] = [] 

200 

201 if not other_args: 

202 # only .2fas file entered - switch to interactive 

203 return command_interactive(filename) 

204 

205 for query in other_args: 

206 found.extend(storage.find(query)) 

207 

208 for twofa in found: 

209 print_for_service(twofa) 

210 

211 

212def get_setting(key: str) -> None: 

213 """ 

214 `--setting key` to get a specifi setting's value. 

215 """ 

216 value = get_cli_setting(key) 

217 rich.print(f"- {key}: {value}") 

218 

219 

220def set_setting(key: str, value: str) -> None: 

221 """ 

222 `--setting key value` to update a setting. 

223 """ 

224 set_cli_setting(key, value) 

225 

226 

227def list_settings() -> None: 

228 """ 

229 Use --settings to show all current settings. 

230 """ 

231 rich.print("Current settings:") 

232 for key, value in state.settings.__dict__.items(): 

233 if key.startswith("_"): 

234 continue 

235 

236 rich.print(f"- {key}: {value}") 

237 

238 

239@clear 

240def set_default_file_interactive(filename: str) -> None: 

241 """ 

242 Interactive menu (after Settings) to set the default 2fas file. 

243 """ 

244 new_filename = questionary.select( 

245 "Pick a file:", 

246 choices=state.settings.files or [], 

247 default=filename, 

248 style=generate_custom_style(), 

249 use_shortcuts=True, 

250 ).ask() 

251 

252 if new_filename is None: 

253 return command_settings(filename) 

254 

255 set_setting("default-file", new_filename) 

256 prepare_to_generate(new_filename) # ask for passphrase 

257 

258 return command_settings(new_filename) 

259 

260 

261@clear() 

262def command_manage_files(filename: str = None) -> None: 

263 """ 

264 Interactive menu (after Settings) to manage known files. 

265 """ 

266 to_remove = questionary.checkbox( 

267 "Which files do you want to remove?", 

268 choices=state.settings.files or [], 

269 style=generate_custom_style(), 

270 ).ask() 

271 if to_remove is not None: 

272 state.settings.remove_file(to_remove) 

273 

274 if filename: 

275 return command_settings(filename) 

276 

277 return None 

278 

279 

280@clear 

281def toggle_autoverbose(filename: str) -> None: 

282 """ 

283 Interactive menu to manage the 'auto verbose' setting. 

284 """ 

285 settings = state.settings 

286 

287 is_enabled = "yes" if settings.auto_verbose else "no" 

288 color = "green" if settings.auto_verbose else "red" 

289 rich.print(f"[blue]Auto Verbose enabled:[/blue] [{color}]{is_enabled}[/{color}]") 

290 

291 text_enabled = "Enable" 

292 new_value = ( 

293 questionary.select( 

294 "Use Auto Verbose?", 

295 choices=[ 

296 text_enabled, 

297 "Disable", 

298 ], 

299 style=generate_custom_style(), 

300 ).ask() 

301 == text_enabled 

302 ) 

303 

304 settings.auto_verbose = new_value 

305 state.verbose = new_value 

306 set_cli_setting("auto_verbose", new_value) 

307 return command_settings(filename) 

308 

309 

310@clear 

311def command_settings(filename: str) -> None: 

312 """ 

313 Menu that shows up when you've chosen 'Settings' from the interactive menu. 

314 """ 

315 rich.print(f"Active file: [blue]{filename}[/blue]") 

316 action = questionary.select( 

317 "What do you want to do?", 

318 choices=generate_choices( 

319 { 

320 "Show current settings": "show-settings", 

321 "Set default file": "set-default-file", 

322 "Add file": "add-file", 

323 "Remove files": "remove-files", 

324 "Toggle auto-verbose": "auto-verbose", 

325 "Back": "back", 

326 } 

327 ), 

328 use_shortcuts=True, 

329 style=generate_custom_style(), 

330 ).ask() 

331 

332 match action: 

333 case "show-settings": 

334 return command_setting([]) 

335 case "set-default-file": 

336 set_default_file_interactive(filename) 

337 case "add-file": 

338 prepare_to_generate(add_2fas_file()) 

339 return command_settings(filename) 

340 case "remove-files": 

341 return command_manage_files(filename) 

342 case "back": 

343 return command_interactive(filename) 

344 case "auto-verbose": 

345 return toggle_autoverbose(filename) 

346 case _: 

347 exit_with_clear(1) 

348 

349 

350def command_setting(args: list[str]) -> None: 

351 """ 

352 Triggered when using --setting, --settings, -s. 

353 

354 Multiple options: 

355 --setting 

356 --setting key 

357 --setting key value, --setting key=value 

358 """ 

359 # required until PyCharm understands 'match' better: 

360 keyvalue: str 

361 key: str 

362 value: str 

363 

364 match args: 

365 case []: 

366 list_settings() 

367 case [keyvalue]: 

368 # key=value 

369 if "=" not in keyvalue: 

370 # get setting 

371 get_setting(keyvalue) 

372 else: 

373 # set settings 

374 set_setting(*keyvalue.split("=", 1)) 

375 case [key, value]: 

376 set_setting(key, value) 

377 case other: 

378 raise ValueError(f"Can't set setting '{other}'.") 

379 

380 

381def command_update() -> None: 

382 """ 

383 --self-update tries to update this library to the latest version on pypi. 

384 """ 

385 python = sys.executable 

386 pip = f"{python} -m pip" 

387 cmd = f"{pip} install --upgrade 2fas" 

388 if os.system(cmd): # nosec: B605 

389 rich.print("[red] could not self-update [/red]") 

390 else: 

391 rich.print("[green] 2fas is at the latest version [/green]") 

392 

393 

394def print_version() -> None: 

395 """ 

396 --version prints the currently installed version of this library. 

397 """ 

398 rich.print(__version__) 

399 

400 

401@app.command() 

402def main( 

403 args: list[str] = typer.Argument(None), 

404 # mutually exclusive actions: 

405 setting: bool = typer.Option( 

406 False, 

407 "--setting", 

408 "--settings", 

409 "-s", 

410 help="Use `--setting` without an argument to see all settings. " 

411 "Use `--setting <name>` to see the current value of a setting. " 

412 "Use `--setting <name> <value>` to update a setting.", 

413 ), 

414 info: str = typer.Option( 

415 None, "--info", "-i", help="`--info <service>` show all known info about a TOTP service from your .2fas file." 

416 ), 

417 self_update: bool = typer.Option( 

418 False, "--self-update", "-u", help="Try to update the 2fas tool to the latest version." 

419 ), 

420 generate_all: bool = typer.Option(False, "--all", "-a", help="Generate all TOTP codes from the active file."), 

421 version: bool = typer.Option(False, "--version", help="Show the current version of the 2fas cli tool."), 

422 remove: bool = typer.Option( 

423 False, "--remove", "--rm", "-r", help="`--remove <filename>` to remove a .2fas file from the known files" 

424 ), 

425 # flags: 

426 verbose: bool = typer.Option( 

427 False, 

428 "--verbose", 

429 "-v", 

430 help="Show more details (e.g. the username for a TOTP service). " 

431 "You can use `auto-verbose` in settings to always show more info.", 

432 ), 

433) -> None: # pragma: no cover 

434 """ 

435 You can use this command in multiple ways. 

436 

437 2fas 

438 

439 2fas path/to/file.fas <service> 

440 

441 2fas <service> path/to/file.fas 

442 

443 2fas <subcommand> 

444 

445 2fas --setting key value 

446 

447 2fas --setting key=value 

448 """ 

449 # stateless actions: 

450 if version: 

451 return print_version() 

452 elif self_update: 

453 return command_update() 

454 

455 # stateful: 

456 

457 settings = load_cli_settings() 

458 state.update(verbose=settings.auto_verbose or verbose, settings=settings) 

459 

460 file_args = [_ for _ in args if _.endswith(".2fas")] 

461 if len(file_args) > 1: 

462 rich.print("[red]Err: can't work on multiple .2fas files![/red]", file=sys.stderr) 

463 exit(1) 

464 

465 filename = expand_path(file_args[0] if file_args else default_2fas_file()) 

466 settings.add_file(filename) 

467 

468 other_args = [_ for _ in args if not _.endswith(".2fas")] 

469 

470 if setting: 

471 command_setting(args) 

472 elif remove and file_args: 

473 settings.remove_file(file_args[0]) 

474 elif remove: 

475 command_manage_files() 

476 elif info: 

477 services = prepare_to_generate(filename) 

478 show_service_info(services, about=info) 

479 elif generate_all: 

480 services = prepare_to_generate(filename) 

481 generate_all_totp(services) 

482 elif args: 

483 command_generate(filename, other_args) 

484 else: 

485 command_interactive(filename)