"""This is the main entry point for the Gemini Key Management command-line tool.""" import argparse import logging import sys import os import concurrent.futures from . import utils, config, auth, database, actions def main(): """Parses command-line arguments and orchestrates the key management process.""" parser = argparse.ArgumentParser(description="Manage Gemini API keys in Google Cloud projects.") parser.add_argument("action", choices=['create', 'delete'], help="The action to perform: 'create' or 'delete' API keys.") parser.add_argument("--email", help="Specify a single email address to process. Required for 'delete'. If not provided for 'create', emails will be read from emails.txt.") parser.add_argument("--dry-run", action="store_true", help="Simulate the run without making any actual changes to Google Cloud resources.") parser.add_argument("--max-workers", type=int, default=5, help="The maximum number of concurrent projects to process.") parser.add_argument("--auth-retries", type=int, default=3, help="Number of retries for a failed authentication attempt.") parser.add_argument("--auth-retry-delay", type=int, default=5, help="Delay in seconds between authentication retries.") args = parser.parse_args() utils.setup_logging() logging.info(f"Program arguments: {vars(args)}") if args.action == 'delete' and not args.email: parser.error("the --email argument is required for the 'delete' action") if not os.path.exists(config.CLIENT_SECRETS_FILE): logging.error(f"OAuth client secrets file not found at '{config.CLIENT_SECRETS_FILE}'") logging.error("Please follow the setup instructions in README.md to create it.") sys.exit(1) if not os.path.exists(config.CREDENTIALS_DIR): os.makedirs(config.CREDENTIALS_DIR) schema = database.load_schema(config.API_KEYS_SCHEMA_FILE) api_keys_data = database.load_keys_database(config.API_KEYS_DATABASE_FILE, schema) emails_to_process = [] if args.email: emails_to_process.append(args.email) elif args.action == 'delete': logging.error("The 'delete' action requires the --email argument to specify which account's keys to delete.") sys.exit(1) else: # 'create' action without a specific email emails_to_process = utils.load_emails_from_file(config.EMAILS_FILE) if not emails_to_process: logging.info("No emails found in emails.txt. Exiting.") sys.exit(1) creds_map = {} emails_needing_interactive_auth = [] logging.info("Checking credentials and refreshing tokens for all accounts...") with concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor: future_to_email = {executor.submit(auth.get_and_refresh_credentials, email, max_retries=args.auth_retries, retry_delay=args.auth_retry_delay): email for email in emails_to_process} for future in concurrent.futures.as_completed(future_to_email): email = future_to_email[future] try: creds = future.result() if creds: creds_map[email] = creds else: emails_needing_interactive_auth.append(email) except Exception as exc: logging.error(f"Credential check for {email} generated an exception: {exc}", exc_info=True) emails_needing_interactive_auth.append(email) if emails_needing_interactive_auth: logging.info(f"\n--- INTERACTIVE AUTHENTICATION REQUIRED ---") logging.info(f"The following accounts require manual authentication: {', '.join(sorted(emails_needing_interactive_auth))}") for email in sorted(emails_needing_interactive_auth): creds = auth.run_interactive_auth(email, max_retries=args.auth_retries, retry_delay=args.auth_retry_delay) if creds: logging.info(f"Successfully authenticated {email}.") creds_map[email] = creds else: logging.warning(f"Authentication failed or was cancelled for {email}. This account will be skipped.") logging.info("\n--- Credential checking complete ---") for email in emails_to_process: if email in creds_map: actions.process_account(email, creds_map[email], args.action, api_keys_data, dry_run=args.dry_run, max_workers=args.max_workers) else: logging.warning(f"Skipping account {email} because authentication was not successful.") if not args.dry_run: database.save_keys_to_json(api_keys_data, config.API_KEYS_DATABASE_FILE, schema)