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

130 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-22 21:51 +0100

1import os 

2import sys 

3import typing 

4 

5import questionary 

6import rich 

7import typer 

8 

9from .__about__ import __version__ 

10from ._security import keyring_manager 

11from ._types import TwoFactorAuthDetails 

12from .cli_settings import get_cli_setting, load_cli_settings, set_cli_setting 

13from .cli_support import clear, exit_with_clear, generate_custom_style, state 

14from .core import TwoFactorStorage, load_services 

15 

16app = typer.Typer() 

17 

18TwoFactorDetailStorage: typing.TypeAlias = TwoFactorStorage[TwoFactorAuthDetails] 

19 

20 

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

22 keyring_manager.cleanup_keyring() 

23 return load_services(filename or default_2fas_file()) 

24 

25 

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

27 service_name = service.name 

28 code = service.generate() 

29 

30 if state.verbose: 

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

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

33 else: 

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

35 

36 

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

38 for service in services: 

39 print_for_service(service) 

40 

41 

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

43 while service_name := questionary.autocomplete( 

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

45 ).ask(): 

46 for service in services.find(service_name): 

47 print_for_service(service) 

48 

49 

50@clear 

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

52 rich.print(services[about]) 

53 

54 

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

56 while about := questionary.select( 

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

58 ).ask(): 

59 show_service_info(services, about) 

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

61 exit_with_clear(0) 

62 

63 

64@clear 

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

66 if not filename: 

67 # get from settings or 

68 filename = default_2fas_file() 

69 

70 services = prepare_to_generate(filename) 

71 

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

73 

74 match questionary.select( 

75 "What do you want to do?", 

76 choices=[ 

77 questionary.Choice("Generate a TOTP code", "generate-one", shortcut_key="1"), 

78 questionary.Choice("Generate all TOTP codes", "generate-all", shortcut_key="2"), 

79 questionary.Choice("Info about a Service", "see-info", shortcut_key="3"), 

80 questionary.Choice("Settings", "settings", shortcut_key="4"), 

81 questionary.Choice("Exit", "exit", shortcut_key="0"), 

82 ], 

83 use_shortcuts=True, 

84 style=generate_custom_style(), 

85 ).ask(): 

86 case "generate-one": 

87 # query list of items 

88 return generate_one_otp(services) 

89 case "generate-all": 

90 # show all 

91 return generate_all_totp(services) 

92 case "see-info": 

93 return show_service_info_interactive(services) 

94 case "settings": 

95 return command_settings(filename) 

96 # manage files 

97 # change specific settings 

98 # default file - choose from list of files 

99 case _: 

100 exit_with_clear(0) 

101 

102 

103def default_2fas_file() -> str: 

104 settings = state.settings 

105 if settings.default_file: 

106 return settings.default_file 

107 

108 elif settings.files: 

109 return settings.files[0] 

110 

111 filename: str = questionary.path( 

112 "Path to .2fas file?", 

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

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

115 style=generate_custom_style(), 

116 ).ask() 

117 

118 set_cli_setting("default-file", filename) 

119 settings.add_file(filename) 

120 

121 return filename 

122 

123 

124def default_2fas_services() -> TwoFactorDetailStorage: 

125 filename = default_2fas_file() 

126 return prepare_to_generate(filename) 

127 

128 

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

130 storage = prepare_to_generate(filename) 

131 found: list[TwoFactorAuthDetails] = [] 

132 

133 if not other_args: 

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

135 return command_interactive(filename) 

136 

137 for query in other_args: 

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

139 

140 for twofa in found: 

141 print_for_service(twofa) 

142 

143 

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

145 value = get_cli_setting(key) 

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

147 

148 

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

150 set_cli_setting(key, value) 

151 

152 

153def list_settings() -> None: 

154 rich.print("Current settings:") 

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

156 if key.startswith("_"): 

157 continue 

158 

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

160 

161 

162@clear 

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

164 new_filename = questionary.select( 

165 "Pick a file:", 

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

167 default=filename, 

168 style=generate_custom_style(), 

169 use_shortcuts=True, 

170 ).ask() 

171 

172 set_setting("default-file", new_filename) 

173 prepare_to_generate(new_filename) # ask for passphrase 

174 

175 return command_settings(new_filename) 

176 

177 

178@clear 

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

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

181 action = questionary.select( 

182 "What do you want to do?", 

183 choices=[ 

184 questionary.Choice("Set default file", "set-default-file", shortcut_key="1"), 

185 questionary.Choice("Manage files", "manage-files", shortcut_key="2"), 

186 questionary.Choice("Back", "back", shortcut_key="3"), 

187 questionary.Choice("Exit", "exit", shortcut_key="0"), 

188 ], 

189 use_shortcuts=True, 

190 style=generate_custom_style(), 

191 ).ask() 

192 

193 match action: 

194 case "set-default-file": 

195 set_default_file_interactive(filename) 

196 case "manage-files": 

197 print("todo: manage files") 

198 case "back": 

199 return command_interactive(filename) 

200 case _: 

201 exit_with_clear(1) 

202 

203 

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

205 # required until PyCharm understands 'match' better: 

206 keyvalue: str 

207 key: str 

208 value: str 

209 

210 match args: 

211 case []: 

212 list_settings() 

213 case [keyvalue]: 

214 # key=value 

215 if "=" not in keyvalue: 

216 # get setting 

217 get_setting(keyvalue) 

218 else: 

219 # set settings 

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

221 case [key, value]: 

222 set_setting(key, value) 

223 case other: 

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

225 

226 

227def command_update() -> None: 

228 python = sys.executable 

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

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

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

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

233 else: 

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

235 

236 

237def print_version(): 

238 rich.print(__version__) 

239 

240 

241@app.command() 

242def main( 

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

244 # mutually exclusive actions: 

245 setting: bool = typer.Option(False, "--setting", "--settings", "-s"), 

246 info: str = typer.Option(None, "--info", "-i"), 

247 self_update: bool = typer.Option(False, "--self-update", "-u"), 

248 generate_all: bool = typer.Option(False, "--all", "-a"), 

249 version: bool = typer.Option(False, "--version"), 

250 

251 # flags: 

252 verbose: bool = typer.Option(False, "--verbose", "-v"), 

253) -> None: # pragma: no cover 

254 """ 

255 Cli entrypoint. 

256 """ 

257 # 2fas 

258 

259 # 2fas path/to/file.fas <service> 

260 # 2fas <service> path/to/file.fas 

261 # 2fas <subcommand> 

262 

263 # 2fas --setting key value 

264 # 2fas --setting key=value 

265 

266 # stateless actions: 

267 if version: 

268 return print_version() 

269 elif self_update: 

270 command_update() 

271 

272 # stateful: 

273 

274 settings = load_cli_settings() 

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

276 

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

278 if len(file_args) > 1: 

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

280 exit(1) 

281 

282 filename = file_args[0] if file_args else default_2fas_file() 

283 settings.add_file(filename) 

284 

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

286 

287 if setting: 

288 command_setting(args) 

289 elif info: 

290 services = prepare_to_generate(filename) 

291 show_service_info(services, about=info) 

292 elif generate_all: 

293 services = prepare_to_generate(filename) 

294 generate_all_totp(services) 

295 elif args: 

296 command_generate(filename, other_args) 

297 else: 

298 command_interactive(filename) 

299 

300 # todo: something to --remove files from history 

301 # todo: better --help info