Mirrowel commited on
Commit
7830a78
·
1 Parent(s): f5ccdf6

refactor(credential-tool): 🔨 add export submenu for credential management

Browse files

Introduced a new submenu for exporting credentials to .env format to improve user experience and code organization.

- Add `export_credentials_submenu()` function to consolidate all export options
- Implement `export_antigravity_to_env()` for Antigravity credential export
- Refactor main menu to replace individual export options (3, 4, 5) with single "Export Credentials" option
- Maintain consistent UI/UX patterns across all export functions
- Generate .env files with metadata headers and timestamp information

This change improves menu navigation by reducing clutter in the main menu and grouping related export functionality together.

Files changed (1) hide show
  1. src/rotator_library/credential_tool.py +140 -17
src/rotator_library/credential_tool.py CHANGED
@@ -533,6 +533,143 @@ async def export_iflow_to_env():
533
  console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error"))
534
 
535
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
536
  async def main(clear_on_start=True):
537
  """
538
  An interactive CLI tool to add new credentials.
@@ -556,9 +693,7 @@ async def main(clear_on_start=True):
556
  Text.from_markup(
557
  "1. Add OAuth Credential\n"
558
  "2. Add API Key\n"
559
- "3. Export Gemini CLI credential to .env\n"
560
- "4. Export Qwen Code credential to .env\n"
561
- "5. Export iFlow credential to .env"
562
  ),
563
  title="Choose credential type",
564
  style="bold blue"
@@ -566,7 +701,7 @@ async def main(clear_on_start=True):
566
 
567
  setup_type = Prompt.ask(
568
  Text.from_markup("[bold]Please select an option or type [red]'q'[/red] to quit[/bold]"),
569
- choices=["1", "2", "3", "4", "5", "q"],
570
  show_choices=False
571
  )
572
 
@@ -622,19 +757,7 @@ async def main(clear_on_start=True):
622
  input()
623
 
624
  elif setup_type == "3":
625
- await export_gemini_cli_to_env()
626
- console.print("\n[dim]Press Enter to return to main menu...[/dim]")
627
- input()
628
-
629
- elif setup_type == "4":
630
- await export_qwen_code_to_env()
631
- console.print("\n[dim]Press Enter to return to main menu...[/dim]")
632
- input()
633
-
634
- elif setup_type == "5":
635
- await export_iflow_to_env()
636
- console.print("\n[dim]Press Enter to return to main menu...[/dim]")
637
- input()
638
 
639
  def run_credential_tool(from_launcher=False):
640
  """
 
533
  console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error"))
534
 
535
 
536
+ async def export_antigravity_to_env():
537
+ """
538
+ Export an Antigravity credential JSON file to .env format.
539
+ Generates one .env file per credential.
540
+ """
541
+ console.print(Panel("[bold cyan]Export Antigravity Credential to .env[/bold cyan]", expand=False))
542
+
543
+ # Find all antigravity credentials
544
+ antigravity_files = list(OAUTH_BASE_DIR.glob("antigravity_oauth_*.json"))
545
+
546
+ if not antigravity_files:
547
+ console.print(Panel("No Antigravity credentials found. Please add one first using 'Add OAuth Credential'.",
548
+ style="bold red", title="No Credentials"))
549
+ return
550
+
551
+ # Display available credentials
552
+ cred_text = Text()
553
+ for i, cred_file in enumerate(antigravity_files):
554
+ try:
555
+ with open(cred_file, 'r') as f:
556
+ creds = json.load(f)
557
+ email = creds.get("_proxy_metadata", {}).get("email", "unknown")
558
+ cred_text.append(f" {i + 1}. {cred_file.name} ({email})\n")
559
+ except Exception as e:
560
+ cred_text.append(f" {i + 1}. {cred_file.name} (error reading: {e})\n")
561
+
562
+ console.print(Panel(cred_text, title="Available Antigravity Credentials", style="bold blue"))
563
+
564
+ choice = Prompt.ask(
565
+ Text.from_markup("[bold]Please select a credential to export or type [red]'b'[/red] to go back[/bold]"),
566
+ choices=[str(i + 1) for i in range(len(antigravity_files))] + ["b"],
567
+ show_choices=False
568
+ )
569
+
570
+ if choice.lower() == 'b':
571
+ return
572
+
573
+ try:
574
+ choice_index = int(choice) - 1
575
+ if 0 <= choice_index < len(antigravity_files):
576
+ cred_file = antigravity_files[choice_index]
577
+
578
+ # Load the credential
579
+ with open(cred_file, 'r') as f:
580
+ creds = json.load(f)
581
+
582
+ # Extract metadata
583
+ email = creds.get("_proxy_metadata", {}).get("email", "unknown")
584
+
585
+ # Generate .env file name
586
+ safe_email = email.replace("@", "_at_").replace(".", "_")
587
+ env_filename = f"antigravity_{safe_email}.env"
588
+ env_filepath = OAUTH_BASE_DIR / env_filename
589
+
590
+ # Build .env content
591
+ env_lines = [
592
+ f"# Antigravity Credential for: {email}",
593
+ f"# Generated from: {cred_file.name}",
594
+ f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}",
595
+ "",
596
+ f"ANTIGRAVITY_ACCESS_TOKEN={creds.get('access_token', '')}",
597
+ f"ANTIGRAVITY_REFRESH_TOKEN={creds.get('refresh_token', '')}",
598
+ f"ANTIGRAVITY_EXPIRY_DATE={creds.get('expiry_date', 0)}",
599
+ f"ANTIGRAVITY_CLIENT_ID={creds.get('client_id', '')}",
600
+ f"ANTIGRAVITY_CLIENT_SECRET={creds.get('client_secret', '')}",
601
+ f"ANTIGRAVITY_TOKEN_URI={creds.get('token_uri', 'https://oauth2.googleapis.com/token')}",
602
+ f"ANTIGRAVITY_UNIVERSE_DOMAIN={creds.get('universe_domain', 'googleapis.com')}",
603
+ f"ANTIGRAVITY_EMAIL={email}",
604
+ ]
605
+
606
+ # Write to .env file
607
+ with open(env_filepath, 'w') as f:
608
+ f.write('\n'.join(env_lines))
609
+
610
+ success_text = Text.from_markup(
611
+ f"Successfully exported credential to [bold yellow]'{env_filepath}'[/bold yellow]\n\n"
612
+ f"To use this credential:\n"
613
+ f"1. Copy [bold yellow]{env_filepath.name}[/bold yellow] to your deployment environment\n"
614
+ f"2. Load the variables: [bold cyan]export $(cat {env_filepath.name} | grep -v '^#' | xargs)[/bold cyan]\n"
615
+ f"3. Or source it: [bold cyan]source {env_filepath.name}[/bold cyan]\n"
616
+ f"4. The Antigravity provider will automatically use these environment variables"
617
+ )
618
+ console.print(Panel(success_text, style="bold green", title="Success"))
619
+ else:
620
+ console.print("[bold red]Invalid choice. Please try again.[/bold red]")
621
+ except ValueError:
622
+ console.print("[bold red]Invalid input. Please enter a number or 'b'.[/bold red]")
623
+ except Exception as e:
624
+ console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error"))
625
+
626
+
627
+ async def export_credentials_submenu():
628
+ """
629
+ Submenu for credential export options.
630
+ """
631
+ while True:
632
+ console.clear()
633
+ console.print(Panel("[bold cyan]Export Credentials to .env[/bold cyan]", title="--- API Key Proxy ---", expand=False))
634
+
635
+ console.print(Panel(
636
+ Text.from_markup(
637
+ "1. Export Gemini CLI credential\n"
638
+ "2. Export Qwen Code credential\n"
639
+ "3. Export iFlow credential\n"
640
+ "4. Export Antigravity credential"
641
+ ),
642
+ title="Choose credential type to export",
643
+ style="bold blue"
644
+ ))
645
+
646
+ export_choice = Prompt.ask(
647
+ Text.from_markup("[bold]Please select an option or type [red]'b'[/red] to go back[/bold]"),
648
+ choices=["1", "2", "3", "4", "b"],
649
+ show_choices=False
650
+ )
651
+
652
+ if export_choice.lower() == 'b':
653
+ break
654
+
655
+ if export_choice == "1":
656
+ await export_gemini_cli_to_env()
657
+ console.print("\n[dim]Press Enter to return to export menu...[/dim]")
658
+ input()
659
+ elif export_choice == "2":
660
+ await export_qwen_code_to_env()
661
+ console.print("\n[dim]Press Enter to return to export menu...[/dim]")
662
+ input()
663
+ elif export_choice == "3":
664
+ await export_iflow_to_env()
665
+ console.print("\n[dim]Press Enter to return to export menu...[/dim]")
666
+ input()
667
+ elif export_choice == "4":
668
+ await export_antigravity_to_env()
669
+ console.print("\n[dim]Press Enter to return to export menu...[/dim]")
670
+ input()
671
+
672
+
673
  async def main(clear_on_start=True):
674
  """
675
  An interactive CLI tool to add new credentials.
 
693
  Text.from_markup(
694
  "1. Add OAuth Credential\n"
695
  "2. Add API Key\n"
696
+ "3. Export Credentials"
 
 
697
  ),
698
  title="Choose credential type",
699
  style="bold blue"
 
701
 
702
  setup_type = Prompt.ask(
703
  Text.from_markup("[bold]Please select an option or type [red]'q'[/red] to quit[/bold]"),
704
+ choices=["1", "2", "3", "q"],
705
  show_choices=False
706
  )
707
 
 
757
  input()
758
 
759
  elif setup_type == "3":
760
+ await export_credentials_submenu()
 
 
 
 
 
 
 
 
 
 
 
 
761
 
762
  def run_credential_tool(from_launcher=False):
763
  """