Coverage for src/clauth/commands/init.py: 75%

1# Copyright (c) 2025 Mahmood Khordoo 

2# 

3# This software is licensed under the MIT License. 

4# See the LICENSE file in the root directory for details. 

5 

6import subprocess 

7import os 

8import typer 

9import clauth.aws_utils as aws 

10from clauth.config import get_config_manager 

11from clauth.helpers import ( 

12 ExecutableNotFoundError, 

13 clear_screen, 

14 get_app_path, 

15 prompt_for_region_if_needed, 

16 show_welcome_logo, 

17 choose_auth_method, 

18) 

19from clauth.aws_utils import ( 

20 setup_sso_auth, 

21 setup_iam_user_auth, 

22) 

23from rich.console import Console 

24from InquirerPy import inquirer 

25from InquirerPy import get_style 

26 

27app = typer.Typer() 

28console = Console() 

29env = os.environ.copy() 

30 

31 

32def _handle_authentication(config, cli_overrides): 

33 """Handles the logic for AWS authentication.""" 

34 auth_method = choose_auth_method() 

35 typer.echo() 

36 

37 if not prompt_for_region_if_needed(config, cli_overrides): 

38 raise typer.Exit(1) 

39 

40 if auth_method == "skip": 

41 if not aws.user_is_authenticated(profile=config.aws.profile): 

42 typer.secho( 

43 "❌ No valid authentication found. Please choose an authentication method.", 

44 fg=typer.colors.RED, 

45 ) 

46 raise typer.Exit(1) 

47 typer.secho( 

48 f"✅ Already authenticated with AWS profile '{config.aws.profile}'", 

49 fg=typer.colors.GREEN, 

50 ) 

51 typer.echo("Skipping credential setup...") 

52 elif auth_method == "iam": 

53 if not setup_iam_user_auth(config.aws.profile, config.aws.region): 

54 raise typer.Exit(1) 

55 elif auth_method == "sso": 

56 if not setup_sso_auth(config, cli_overrides): 

57 raise typer.Exit(1) 

58 

59 

60def _launch_claude_cli(config, env): 

61 """Launches the Claude Code CLI with the given environment.""" 

62 try: 

63 claude_path = get_app_path(config.cli.claude_cli_name) 

64 clear_screen() 

65 subprocess.run([claude_path], env=env, check=True) 

66 except ExecutableNotFoundError as e: 

67 typer.secho(f"Setup failed: {e}", fg=typer.colors.RED) 

68 typer.secho( 

69 "Please install Claude Code CLI and ensure it's in your PATH.", 

70 fg=typer.colors.YELLOW, 

71 ) 

72 raise typer.Exit(1) 

73 except ValueError as e: 

74 typer.secho(f"Configuration error: {e}", fg=typer.colors.RED) 

75 raise typer.Exit(1) 

76 

77 

78def _handle_model_selection(config, config_manager, console): 

79 """Handles the logic for selecting Bedrock models.""" 

80 # Check if we have existing model configuration 

81 if config.models.default_model_arn and config.models.fast_model_arn: 

82 typer.echo("Found existing model configuration:") 

83 typer.echo(f" Default model: {config.models.default_model}") 

84 typer.echo(f" Small/Fast model: {config.models.fast_model}") 

85 

86 # Get custom style from config manager 

87 custom_style = get_style(config_manager.get_custom_style()) 

88 

89 use_existing = inquirer.confirm( 

90 message="Use existing model configuration?", 

91 default=True, 

92 style=custom_style, 

93 ).execute() 

94 

95 if use_existing: 

96 model_id_default = config.models.default_model 

97 model_id_fast = config.models.fast_model 

98 model_map = { 

99 model_id_default: config.models.default_model_arn, 

100 model_id_fast: config.models.fast_model_arn, 

101 } 

102 typer.echo(f"Using saved models: {model_id_default}, {model_id_fast}") 

103 return model_id_default, model_id_fast, model_map 

104 

105 # No existing configuration or user chose not to use it, do full model discovery 

106 with console.status("[bold blue]Discovering available models...") as status: 

107 model_ids, model_arns = aws.list_bedrock_profiles( 

108 profile=config.aws.profile, 

109 region=config.aws.region, 

110 provider=config.models.provider_filter, 

111 ) 

112 

113 # Get custom style from config manager 

114 custom_style = get_style(config_manager.get_custom_style()) 

115 

116 model_id_default = inquirer.select( 

117 message="Select your [default] model:", 

118 instruction="↑↓ move • Enter select", 

119 pointer="▶ ", 

120 amark="✔", 

121 choices=model_ids, 

122 default=config.models.default_model 

123 if config.models.default_model in model_ids 

124 else (model_ids[0] if model_ids else None), 

125 style=custom_style, 

126 max_height="100%", 

127 ).execute() 

128 

129 model_id_fast = inquirer.select( 

130 message="Select your [small/fast] model (you can choose the same as default):", 

131 instruction="↑↓ move • Enter select", 

132 pointer="▶ ", 

133 amark="✔", 

134 choices=model_ids, 

135 default=config.models.fast_model 

136 if config.models.fast_model in model_ids 

137 else (model_ids[-1] if model_ids else None), 

138 style=custom_style, 

139 max_height="100%", 

140 ).execute() 

141 

142 model_map = {id: arn for id, arn in zip(model_ids, model_arns)} 

143 

144 # Save updated model selections to configuration 

145 config_manager.update_model_settings( 

146 default_model=model_id_default, 

147 fast_model=model_id_fast, 

148 default_arn=model_map[model_id_default], 

149 fast_arn=model_map[model_id_fast], 

150 ) 

151 return model_id_default, model_id_fast, model_map 

152 

153 

154def init_command( 

155 profile: str = typer.Option( 

156 None, 

157 "--profile", 

158 "-p", 

159 help="AWS profile to create or update (saved under [profile <name>] in ~/.aws/config).", 

160 rich_help_panel="AWS Profile", 

161 ), 

162 session_name: str = typer.Option( 

163 None, 

164 "--session-name", 

165 "-s", 

166 help="Name of the SSO session to create (saved under [sso-session <name>] in ~/.aws/config).", 

167 rich_help_panel="AWS SSO", 

168 ), 

169 sso_start_url: str = typer.Option( 

170 None, 

171 "--sso-start-url", 

172 help="IAM Identity Center (SSO) Start URL (e.g., https://d-…awsapps.com/start/).", 

173 rich_help_panel="AWS SSO", 

174 ), 

175 sso_region: str = typer.Option( 

176 None, 

177 "--sso-region", 

178 help="Region that hosts your IAM Identity Center (SSO) instance.", 

179 rich_help_panel="AWS SSO", 

180 ), 

181 region: str = typer.Option( 

182 None, 

183 "--region", 

184 "-r", 

185 help="Default AWS client region for this profile (used for STS/Bedrock calls).", 

186 rich_help_panel="AWS Profile", 

187 ), 

188 auto_start: bool = typer.Option( 

189 None, 

190 "--auto-start/--no-auto-start", 

191 help="Launch the Claude CLI immediately after successful setup.", 

192 rich_help_panel="Behavior", 

193 ), 

194): 

195 """ 

196 Interactive setup wizard for CLAUTH. 

197 

198 Configures AWS authentication (SSO or IAM user), discovers available Bedrock models, 

199 and optionally launches Claude Code CLI with proper environment variables. 

200 This is the main entry point for first-time CLAUTH setup. 

201 

202 Args: 

203 profile: AWS profile name to create/update (default from config) 

204 session_name: SSO session name (default from config, SSO only) 

205 sso_start_url: IAM Identity Center start URL (default from config, SSO only) 

206 sso_region: SSO region (default from config, SSO only) 

207 region: Default AWS region for profile (default from config) 

208 auto_start: Whether to launch Claude Code after setup (default from config) 

209 """ 

210 # Load configuration and apply CLI overrides 

211 config_manager = get_config_manager() 

212 config = config_manager.load() 

213 

214 # Track which CLI parameters were provided 

215 cli_overrides = { 

216 "profile": profile is not None, 

217 "session_name": session_name is not None, 

218 "sso_start_url": sso_start_url is not None, 

219 "sso_region": sso_region is not None, 

220 "region": region is not None, 

221 "auto_start": auto_start is not None, 

222 } 

223 

224 # Override config with CLI parameters if provided 

225 if profile is not None: 

226 config.aws.profile = profile 

227 if session_name is not None: 

228 config.aws.session_name = session_name 

229 if sso_start_url is not None: 

230 config.aws.sso_start_url = sso_start_url 

231 if sso_region is not None: 

232 config.aws.sso_region = sso_region 

233 if region is not None: 

234 config.aws.region = region 

235 if auto_start is not None: 

236 config.cli.auto_start = auto_start 

237 

238 show_welcome_logo(console=console) 

239 

240 try: 

241 typer.secho( 

242 "Step 1/3 — Configuring AWS authentication...", fg=typer.colors.BLUE 

243 ) 

244 typer.echo() 

245 

246 _handle_authentication(config, cli_overrides) 

247 

248 typer.secho("Step 2/3 — Configuring models...", fg=typer.colors.BLUE) 

249 

250 model_id_default, model_id_fast, model_map = _handle_model_selection(config, config_manager, console) 

251 

252 typer.echo(f"Default model: {model_id_default}") 

253 typer.echo(f"Small/Fast model: {model_id_fast}") 

254 

255 env.update( 

256 { 

257 "AWS_PROFILE": config.aws.profile, 

258 "AWS_REGION": config.aws.region, 

259 "CLAUDE_CODE_USE_BEDROCK": "1", 

260 "ANTHROPIC_MODEL": model_map[model_id_default], 

261 "ANTHROPIC_SMALL_FAST_MODEL_AWS_REGION": model_map[model_id_fast], 

262 } 

263 ) 

264 

265 typer.echo( 

266 f"default model: {model_id_default}\n small/fast model: {model_id_fast}\n" 

267 ) 

268 

269 if config.cli.auto_start: 

270 typer.secho("Setup complete ✅", fg=typer.colors.GREEN) 

271 typer.secho("Step 3/3 — Launching Claude Code...", fg=typer.colors.BLUE) 

272 _launch_claude_cli(config, env) 

273 else: 

274 typer.secho("Step 3/3 — Setup complete.", fg=typer.colors.GREEN) 

275 typer.echo("Run the Claude Code CLI when you're ready: ", nl=False) 

276 typer.secho(config.cli.claude_cli_name, bold=True) 

277 

278 except subprocess.CalledProcessError as e: 

279 typer.secho(f"Setup failed. Exit code: {e.returncode}", fg=typer.colors.RED) 

280 exit(f"Failed to setup. Error Code: {e.returncode}")