disallow >1 api keys per project

This commit is contained in:
2025-08-18 15:23:11 +05:30
parent 0721e8f960
commit 1523a0c600

39
main.py
View File

@@ -1,6 +1,7 @@
import os import os
import sys import sys
import argparse import argparse
import re
import google.auth import google.auth
from google.oauth2.credentials import Credentials from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow from google_auth_oauthlib.flow import InstalledAppFlow
@@ -28,6 +29,18 @@ def load_emails_from_file(filename):
# Ignore empty lines and lines starting with # # Ignore empty lines and lines starting with #
return [line.strip() for line in f if line.strip() and not line.startswith("#")] return [line.strip() for line in f if line.strip() and not line.startswith("#")]
def load_existing_api_keys(email):
"""Loads API keys from the user's key file."""
filename = f"{email}.keys.txt"
keys = set()
if os.path.exists(filename):
with open(filename, "r") as f:
content = f.read()
# Find all occurrences of "Key: <key_string>"
found_keys = re.findall(r"Key:\s*(.+)", content)
keys.update(found_keys)
return keys
def main(): def main():
"""Main function to orchestrate API key creation or deletion.""" """Main function to orchestrate API key creation or deletion."""
parser = argparse.ArgumentParser(description="Manage Gemini API keys in Google Cloud projects.") parser = argparse.ArgumentParser(description="Manage Gemini API keys in Google Cloud projects.")
@@ -66,6 +79,13 @@ def process_account(email, action):
print(f"Could not get credentials for {email}. Skipping.") print(f"Could not get credentials for {email}. Skipping.")
return return
# Load existing keys if we are creating
existing_keys = set()
if action == 'create':
existing_keys = load_existing_api_keys(email)
if existing_keys:
print(f"Loaded {len(existing_keys)} existing keys from file.")
try: try:
resource_manager = resourcemanager_v3.ProjectsClient(credentials=creds) resource_manager = resourcemanager_v3.ProjectsClient(credentials=creds)
projects = list(resource_manager.search_projects()) projects = list(resource_manager.search_projects())
@@ -80,6 +100,11 @@ def process_account(email, action):
print(f"- Project: {project_id}") print(f"- Project: {project_id}")
if action == 'create': if action == 'create':
# Check if a key with the specific display name already exists
if project_has_gemini_key(project_id, creds):
print(" 'Gemini API Key' already exists in this project. Skipping creation.")
continue
if enable_api(project_id, creds): if enable_api(project_id, creds):
key = create_api_key(project_id, creds) key = create_api_key(project_id, creds)
if key: if key:
@@ -93,6 +118,20 @@ def process_account(email, action):
except google_exceptions.GoogleAPICallError as err: except google_exceptions.GoogleAPICallError as err:
print(f"An API error occurred while processing account {email}: {err}") print(f"An API error occurred while processing account {email}: {err}")
def project_has_gemini_key(project_id, credentials):
"""Checks if a project already has a key named '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)
for key in keys:
if key.display_name == "Gemini API Key":
return True
return False
except google_exceptions.GoogleAPICallError as err:
print(f" Could not list keys in project {project_id}. Assuming no key exists. Error: {err}")
return False
def get_credentials_for_email(email): def get_credentials_for_email(email):
"""Handles the OAuth2 flow for a given email.""" """Handles the OAuth2 flow for a given email."""