allow deletion of keys

main.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import argparse
 import google.auth
 from google.oauth2.credentials import Credentials
 from google_auth_oauthlib.flow import InstalledAppFlow
@@ -18,16 +19,22 @@ SCOPES = [
 ]
 
 def load_emails_from_file(filename):
-    """Loads a list of emails from a text file."""
+    """Loads a list of emails from a text file, ignoring comments."""
     if not os.path.exists(filename):
         print(f"Error: Email file not found at '{filename}'")
         print("Please create it and add one email address per line.")
         return []
     with open(filename, "r") as f:
+        # Ignore empty lines and lines starting with #
         return [line.strip() for line in f if line.strip() and not line.startswith("#")]
 
 def main():
-    """Main function to orchestrate API key creation."""
+    """Main function to orchestrate API key creation or deletion."""
+    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' action. If not provided for 'create', emails will be read from emails.txt.")
+    args = parser.parse_args()
+
     if not os.path.exists(CLIENT_SECRETS_FILE):
         print(f"Error: OAuth client secrets file not found at '{CLIENT_SECRETS_FILE}'")
         print("Please follow the setup instructions in README.md to create it.")
@@ -36,39 +43,56 @@ def main():
     if not os.path.exists(CREDENTIALS_DIR):
         os.makedirs(CREDENTIALS_DIR)
 
-    emails = load_emails_from_file(EMAILS_FILE)
-    if not emails:
+    emails_to_process = []
+    if args.email:
+        emails_to_process.append(args.email)
+    elif args.action == 'delete':
+        print("Error: The 'delete' action requires the --email argument.")
         sys.exit(1)
+    else: # action is 'create' and no email provided
+        emails_to_process = load_emails_from_file(EMAILS_FILE)
+        if not emails_to_process:
+            print("No emails found in emails.txt. Exiting.")
+            sys.exit(1)
 
-    for email in emails:
-        print(f"--- Processing account: {email} ---")
-        creds = get_credentials_for_email(email)
-        if not creds:
-            print(f"Could not get credentials for {email}. Skipping.")
-            continue
+    for email in emails_to_process:
+        process_account(email, args.action)
 
-        try:
-            resource_manager = resourcemanager_v3.ProjectsClient(credentials=creds)
-            projects = list(resource_manager.search_projects()) # Convert iterator to list
+def process_account(email, action):
+    """Processes a single account for the given action."""
+    print(f"--- Processing account: {email} for action: {action} ---")
+    creds = get_credentials_for_email(email)
+    if not creds:
+        print(f"Could not get credentials for {email}. Skipping.")
+        return
 
-            if not projects:
-                print(f"No projects found for {email}.")
-                continue
+    try:
+        resource_manager = resourcemanager_v3.ProjectsClient(credentials=creds)
+        projects = list(resource_manager.search_projects())
 
-            print(f"Found {len(projects)} projects. Processing...")
-            for project in projects:
-                project_id = project.project_id
-                print(f"- Project: {project_id}")
+        if not projects:
+            print(f"No projects found for {email}.")
+            return
+
+        print(f"Found {len(projects)} projects. Processing...")
+        for project in projects:
+            project_id = project.project_id
+            print(f"- Project: {project_id}")
+
+            if action == 'create':
                 if enable_api(project_id, creds):
                     key = create_api_key(project_id, creds)
                     if key:
                         save_api_key(email, project_id, key.key_string)
+            elif action == 'delete':
+                delete_api_keys(project_id, creds)
+
+    except google_exceptions.PermissionDenied as err:
+        print(f"Permission denied for account {email}. Check IAM roles.")
+        print(f"  Error: {err}")
+    except google_exceptions.GoogleAPICallError as err:
+        print(f"An API error occurred while processing account {email}: {err}")
 
-        except google_exceptions.PermissionDenied as err:
-            print(f"Permission denied for account {email}. Check IAM roles.")
-            print(f"  Error: {err}")
-        except google_exceptions.GoogleAPICallError as err:
-            print(f"An API error occurred while processing account {email}: {err}")
 
 def get_credentials_for_email(email):
     """Handles the OAuth2 flow for a given email."""
@@ -153,6 +177,35 @@ def create_api_key(project_id, credentials):
         print(f"Error creating API key for project {project_id}: {err}")
         return None
 
+def delete_api_keys(project_id, credentials):
+    """Deletes all API keys with the display name 'Gemini API Key'."""
+    try:
+        api_keys_client = api_keys_v2.ApiKeysClient(credentials=credentials)
+        parent = f"projects/{project_id}/locations/global"
+
+        keys = api_keys_client.list_keys(parent=parent)
+        keys_to_delete = [key for key in keys if key.display_name == "Gemini API Key"]
+
+        if not keys_to_delete:
+            print(f"  No 'Gemini API Key' found to delete.")
+            return
+
+        print(f"  Found {len(keys_to_delete)} key(s) with display name 'Gemini API Key'. Deleting...")
+        for key in keys_to_delete:
+            try:
+                request = api_keys_v2.DeleteKeyRequest(name=key.name)
+                operation = api_keys_client.delete_key(request=request)
+                operation.result()  # Wait for completion
+                print(f"  Successfully deleted key: {key.uid}")
+            except google_exceptions.GoogleAPICallError as err:
+                print(f"  Error deleting key {key.uid}: {err}")
+
+    except google_exceptions.PermissionDenied as err:
+        print(f"  Permission denied to list or delete API keys for project {project_id}. Skipping.")
+    except google_exceptions.GoogleAPICallError as err:
+        print(f"  An API error occurred while deleting keys for project {project_id}: {err}")
+
+
 def save_api_key(email, project_id, api_key):
     """Saves the API key to a file, including the project ID."""
     filename = f"{email}.keys.txt"
@@ -161,4 +214,4 @@ def save_api_key(email, project_id, api_key):
     print(f"Saved API key to {filename}")
 
 if __name__ == "__main__":
-    main()
+    main()