lucky/GeminiKeyManagement
1
0
Fork 0
mirror of https://github.com/not-lucky/GeminiKeyManagement.git synced 2025-12-06 08:44:01 +05:30
Code Issues Packages Projects Releases Wiki Activity

added dry-run feature, logging

Browse Source
This commit is contained in:
not-lucky
2025-08-19 20:35:57 +05:30
parent 144cedf1c7
commit 769cb7b311
5 changed files with 120 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -11,3 +11,4 @@ credentials.json
*.key *.key
emails.txt emails.txt
*.keys.txt *.keys.txt
api_keys_database.json

14
README.md
View File

@@ -29,6 +29,15 @@ This script automates the process of creating and deleting Google Cloud API keys
Run the script from your terminal with a specified action (`create` or `delete`). Run the script from your terminal with a specified action (`create` or `delete`).
### Dry Run
You can simulate any action without making real changes by adding the `--dry-run` flag. This is useful for testing and verifying the script's behavior.
```bash
python main.py create --dry-run
python main.py delete --email your.email@example.com --dry-run
```
### Creating API Keys ### Creating API Keys
- **For all users in `emails.txt`:** - **For all users in `emails.txt`:**
@@ -47,14 +56,15 @@ The script will then:
- Find all Google Cloud projects accessible by the authenticated user. - Find all Google Cloud projects accessible by the authenticated user.
- Enable the `generativelanguage.googleapis.com` API for each project. - Enable the `generativelanguage.googleapis.com` API for each project.
- Create a new API key named "Gemini API Key" with restrictions to the Generative Language API. - Create a new API key named "Gemini API Key" with restrictions to the Generative Language API.
- Save the API key(s) to a file named `<email>.keys.txt`. - Save the API key(s) to a central `api_keys_database.json` file.
### Deleting API Keys ### Deleting API Keys
To delete all API keys with the display name "Gemini API Key" for a specific user, run: To delete all API keys with the display name "Gemini API Key" for a specific user, the `--email` argument is required:
```bash ```bash
python main.py delete --email your.email@example.com python main.py delete --email your.email@example.com
``` ```
This will iterate through all projects accessible by that user and remove the matching keys. This will iterate through all projects accessible by that user and remove the matching keys.

155
main.py
View File

@@ -2,6 +2,7 @@ import os
import sys import sys
import argparse import argparse
import json import json
import logging
from datetime import datetime, timezone from datetime import datetime, timezone
import jsonschema import jsonschema
import google.auth import google.auth
@@ -19,6 +20,15 @@ API_KEYS_DATABASE_FILE = "api_keys_database.json"
API_KEYS_SCHEMA_FILE = os.path.join("schemas", "v1", "api_keys_database.schema.json") API_KEYS_SCHEMA_FILE = os.path.join("schemas", "v1", "api_keys_database.schema.json")
# --------------------- # ---------------------
# --- LOGGING SETUP ---
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
stream=sys.stdout,
)
# ---------------------
SCOPES = [ SCOPES = [
"https://www.googleapis.com/auth/cloud-platform", "https://www.googleapis.com/auth/cloud-platform",
] ]
@@ -26,20 +36,20 @@ SCOPES = [
def load_schema(filename): def load_schema(filename):
"""Loads a JSON schema from a file.""" """Loads a JSON schema from a file."""
if not os.path.exists(filename): if not os.path.exists(filename):
print(f"Error: Schema file not found at '{filename}'") logging.error(f"Schema file not found at '{filename}'")
sys.exit(1) sys.exit(1)
with open(filename, "r") as f: with open(filename, "r") as f:
try: try:
return json.load(f) return json.load(f)
except json.JSONDecodeError: except json.JSONDecodeError:
print(f"Error: Could not decode JSON schema from {filename}.") logging.error(f"Could not decode JSON schema from {filename}.")
sys.exit(1) sys.exit(1)
def load_emails_from_file(filename): def load_emails_from_file(filename):
"""Loads a list of emails from a text file, ignoring comments.""" """Loads a list of emails from a text file, ignoring comments."""
if not os.path.exists(filename): if not os.path.exists(filename):
print(f"Error: Email file not found at '{filename}'") logging.error(f"Email file not found at '{filename}'")
print("Please create it and add one email address per line.") logging.info("Please create it and add one email address per line.")
return [] return []
with open(filename, "r") as f: with open(filename, "r") as f:
# Ignore empty lines and lines starting with # # Ignore empty lines and lines starting with #
@@ -58,9 +68,9 @@ def load_keys_database(filename, schema):
jsonschema.validate(instance=data, schema=schema) jsonschema.validate(instance=data, schema=schema)
return data return data
except json.JSONDecodeError: except json.JSONDecodeError:
print(f"Warning: Could not decode JSON from {filename}. Starting fresh.") logging.warning(f"Could not decode JSON from {filename}. Starting fresh.")
except jsonschema.ValidationError as e: except jsonschema.ValidationError as e:
print(f"Warning: Database file '{filename}' is not valid. {e.message}. Starting fresh.") logging.warning(f"Database file '{filename}' is not valid. {e.message}. Starting fresh.")
return { return {
"schema_version": "1.0.0", "schema_version": "1.0.0",
@@ -70,15 +80,17 @@ def load_keys_database(filename, schema):
def save_keys_to_json(data, filename, schema): def save_keys_to_json(data, filename, schema):
"""Validates and saves the API key data to a single JSON file.""" """Validates and saves the API key data to a single JSON file."""
data["generation_timestamp_utc"] = datetime.now(timezone.utc).isoformat() now = datetime.now(timezone.utc).isoformat()
data["generation_timestamp_utc"] = data.get("generation_timestamp_utc", now)
data["last_modified_utc"] = now
try: try:
jsonschema.validate(instance=data, schema=schema) jsonschema.validate(instance=data, schema=schema)
with open(filename, "w") as f: with open(filename, "w") as f:
json.dump(data, f, indent=2) json.dump(data, f, indent=2)
print(f"--- Database saved to {filename} ---") logging.info(f"--- Database saved to {filename} ---")
except jsonschema.ValidationError as e: except jsonschema.ValidationError as e:
print(f"Error: Data to be saved is invalid. Could not write to '{filename}'.") logging.error(f"Data to be saved is invalid. Could not write to '{filename}'.")
print(f"Validation Error: {e.message}") logging.error(f"Validation Error: {e.message}")
sys.exit(1) sys.exit(1)
@@ -86,12 +98,16 @@ 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.")
parser.add_argument("action", choices=['create', 'delete'], help="The action to perform: 'create' or 'delete' API keys.") 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. If not provided for 'create', emails will be read from emails.txt.") 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.")
args = parser.parse_args() args = parser.parse_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(CLIENT_SECRETS_FILE): if not os.path.exists(CLIENT_SECRETS_FILE):
print(f"Error: OAuth client secrets file not found at '{CLIENT_SECRETS_FILE}'") logging.error(f"OAuth client secrets file not found at '{CLIENT_SECRETS_FILE}'")
print("Please follow the setup instructions in README.md to create it.") logging.error("Please follow the setup instructions in README.md to create it.")
sys.exit(1) sys.exit(1)
if not os.path.exists(CREDENTIALS_DIR): if not os.path.exists(CREDENTIALS_DIR):
@@ -104,25 +120,29 @@ def main():
if args.email: if args.email:
emails_to_process.append(args.email) emails_to_process.append(args.email)
elif args.action == 'delete': elif args.action == 'delete':
print("Error: The 'delete' action requires the --email argument to specify which account's keys to delete.") logging.error("The 'delete' action requires the --email argument to specify which account's keys to delete.")
sys.exit(1) sys.exit(1)
else: # action is 'create' and no email provided else: # action is 'create' and no email provided
emails_to_process = load_emails_from_file(EMAILS_FILE) emails_to_process = load_emails_from_file(EMAILS_FILE)
if not emails_to_process: if not emails_to_process:
print("No emails found in emails.txt. Exiting.") logging.info("No emails found in emails.txt. Exiting.")
sys.exit(1) sys.exit(1)
for email in emails_to_process: for email in emails_to_process:
process_account(email, args.action, api_keys_data) process_account(email, args.action, api_keys_data, args.dry_run)
if not args.dry_run:
save_keys_to_json(api_keys_data, API_KEYS_DATABASE_FILE, schema) save_keys_to_json(api_keys_data, API_KEYS_DATABASE_FILE, schema)
def process_account(email, action, api_keys_data): def process_account(email, action, api_keys_data, dry_run=False):
"""Processes a single account for the given action.""" """Processes a single account for the given action."""
print(f"--- Processing account: {email} for action: {action} ---") logging.info(f"--- Processing account: {email} for action: {action} ---")
if dry_run:
logging.info("*** DRY RUN MODE ENABLED ***")
creds = get_credentials_for_email(email) creds = get_credentials_for_email(email)
if not creds: if not creds:
print(f"Could not get credentials for {email}. Skipping.") logging.warning(f"Could not get credentials for {email}. Skipping.")
return return
account_entry = next((acc for acc in api_keys_data["accounts"] if acc.get("account_details", {}).get("email") == email), None) account_entry = next((acc for acc in api_keys_data["accounts"] if acc.get("account_details", {}).get("email") == email), None)
@@ -144,33 +164,33 @@ def process_account(email, action, api_keys_data):
projects = list(resource_manager.search_projects()) projects = list(resource_manager.search_projects())
if not projects: if not projects:
print(f"No projects found for {email}.") logging.info(f"No projects found for {email}.")
return return
print(f"Found {len(projects)} projects. Processing...") logging.info(f"Found {len(projects)} projects. Processing...")
for project in projects: for project in projects:
project_id = project.project_id project_id = project.project_id
print(f"- Project: {project_id} ({project.display_name})") logging.info(f"- Project: {project_id} ({project.display_name})")
if action == 'create': if action == 'create':
if project_has_gemini_key(project_id, creds): if project_has_gemini_key(project_id, creds):
print(" 'Gemini API Key' already exists in this project. Skipping creation.") logging.info(" 'Gemini API Key' already exists in this project. Skipping creation.")
continue continue
if enable_api(project_id, creds): if enable_api(project_id, creds, dry_run=dry_run):
key_object = create_api_key(project_id, creds) key_object = create_api_key(project_id, creds, dry_run=dry_run)
if key_object: if key_object:
add_key_to_database(account_entry, project, key_object) add_key_to_database(account_entry, project, key_object)
elif action == 'delete': elif action == 'delete':
deleted_keys_uids = delete_api_keys(project_id, creds) deleted_keys_uids = delete_api_keys(project_id, creds, dry_run=dry_run)
if deleted_keys_uids: if deleted_keys_uids:
remove_keys_from_database(account_entry, project_id, deleted_keys_uids) remove_keys_from_database(account_entry, project_id, deleted_keys_uids)
except google_exceptions.PermissionDenied as err: except google_exceptions.PermissionDenied as err:
print(f"Permission denied for account {email}. Check IAM roles.") logging.error(f"Permission denied for account {email}. Check IAM roles.")
print(f" Error: {err}") logging.error(f" Error: {err}")
except google_exceptions.GoogleAPICallError as err: except google_exceptions.GoogleAPICallError as err:
print(f"An API error occurred while processing account {email}: {err}") logging.error(f"An API error occurred while processing account {email}: {err}")
def add_key_to_database(account_entry, project, key_object): def add_key_to_database(account_entry, project, key_object):
"""Adds a new API key's details to the data structure.""" """Adds a new API key's details to the data structure."""
@@ -212,9 +232,9 @@ def add_key_to_database(account_entry, project, key_object):
existing_key = next((k for k in project_entry["api_keys"] if k.get("key_details", {}).get("key_id") == key_object.uid), None) existing_key = next((k for k in project_entry["api_keys"] if k.get("key_details", {}).get("key_id") == key_object.uid), None)
if not existing_key: if not existing_key:
project_entry["api_keys"].append(new_key_entry) project_entry["api_keys"].append(new_key_entry)
print(f" Added key {key_object.uid} to local database for project {project_id}") logging.info(f" Added key {key_object.uid} to local database for project {project_id}")
else: else:
print(f" Key {key_object.uid} already exists in local database for project {project_id}") logging.warning(f" Key {key_object.uid} already exists in local database for project {project_id}")
def remove_keys_from_database(account_entry, project_id, deleted_keys_uids): def remove_keys_from_database(account_entry, project_id, deleted_keys_uids):
"""Removes deleted API keys from the data structure.""" """Removes deleted API keys from the data structure."""
@@ -230,7 +250,7 @@ def remove_keys_from_database(account_entry, project_id, deleted_keys_uids):
final_key_count = len(project_entry["api_keys"]) final_key_count = len(project_entry["api_keys"])
num_removed = initial_key_count - final_key_count num_removed = initial_key_count - final_key_count
if num_removed > 0: if num_removed > 0:
print(f" Removed {num_removed} key(s) from local database for project {project_id}") logging.info(f" Removed {num_removed} key(s) from local database for project {project_id}")
def project_has_gemini_key(project_id, credentials): def project_has_gemini_key(project_id, credentials):
"""Checks if a project already has a key named 'Gemini API Key'.""" """Checks if a project already has a key named 'Gemini API Key'."""
@@ -243,7 +263,7 @@ def project_has_gemini_key(project_id, credentials):
return True return True
return False return False
except google_exceptions.GoogleAPICallError as err: except google_exceptions.GoogleAPICallError as err:
print(f" Could not list keys in project {project_id}. Error: {err}") logging.error(f" Could not list keys in project {project_id}. Error: {err}")
return False return False
def get_credentials_for_email(email): def get_credentials_for_email(email):
@@ -256,10 +276,10 @@ def get_credentials_for_email(email):
if not creds or not creds.valid: if not creds or not creds.valid:
if creds and creds.expired and creds.refresh_token: if creds and creds.expired and creds.refresh_token:
print("Refreshing credentials...") logging.info("Refreshing credentials...")
creds.refresh(google.auth.transport.requests.Request()) creds.refresh(google.auth.transport.requests.Request())
else: else:
print(f"Please authenticate with: {email}") logging.info(f"Please authenticate with: {email}")
flow = InstalledAppFlow.from_client_secrets_file( flow = InstalledAppFlow.from_client_secrets_file(
CLIENT_SECRETS_FILE, SCOPES CLIENT_SECRETS_FILE, SCOPES
) )
@@ -268,10 +288,10 @@ def get_credentials_for_email(email):
token.write(creds.to_json()) token.write(creds.to_json())
return creds return creds
except Exception as e: except Exception as e:
print(f"An unexpected error occurred during authentication: {e}") logging.error(f"An unexpected error occurred during authentication: {e}")
return None return None
def enable_api(project_id, credentials): def enable_api(project_id, credentials, dry_run=False):
"""Enables the Generative Language API if it's not already enabled.""" """Enables the Generative Language API if it's not already enabled."""
service_name = "generativelanguage.googleapis.com" service_name = "generativelanguage.googleapis.com"
service_path = f"projects/{project_id}/services/{service_name}" service_path = f"projects/{project_id}/services/{service_name}"
@@ -282,26 +302,47 @@ def enable_api(project_id, credentials):
service = service_usage_client.get_service(request=service_request) service = service_usage_client.get_service(request=service_request)
if service.state == service_usage_v1.State.ENABLED: if service.state == service_usage_v1.State.ENABLED:
print(f" Generative Language API is already enabled for project {project_id}") logging.info(f" Generative Language API is already enabled for project {project_id}")
return True
logging.info(f" API is not enabled. Attempting to enable...")
if dry_run:
logging.info(f" [DRY RUN] Would enable API for project {project_id}")
return True return True
print(f" API is not enabled. Attempting to enable...")
enable_request = service_usage_v1.EnableServiceRequest(name=service_path) enable_request = service_usage_v1.EnableServiceRequest(name=service_path)
operation = service_usage_client.enable_service(request=enable_request) operation = service_usage_client.enable_service(request=enable_request)
print(" Waiting for API enablement to complete...") logging.info(" Waiting for API enablement to complete...")
operation.result() operation.result()
print(f" Successfully enabled Generative Language API for project {project_id}") logging.info(f" Successfully enabled Generative Language API for project {project_id}")
return True return True
except google_exceptions.PermissionDenied: except google_exceptions.PermissionDenied:
print(f" Permission denied to check or enable API for project {project_id}. Skipping.") logging.warning(f" Permission denied to check or enable API for project {project_id}. Skipping.")
return False return False
except google_exceptions.GoogleAPICallError as err: except google_exceptions.GoogleAPICallError as err:
print(f" Error checking or enabling API for project {project_id}: {err}") logging.error(f" Error checking or enabling API for project {project_id}: {err}")
return False return False
def create_api_key(project_id, credentials): def create_api_key(project_id, credentials, dry_run=False):
"""Creates a new, restricted API key.""" """Creates a new, restricted API key."""
if dry_run:
logging.info(f" [DRY RUN] Would create API key for project {project_id}")
# Return a mock key object for dry run
return api_keys_v2.Key(
name=f"projects/{project_id}/locations/global/keys/mock-key-id",
uid="mock-key-id",
display_name="Gemini API Key",
key_string="mock-key-string-for-dry-run",
create_time=datetime.now(timezone.utc),
update_time=datetime.now(timezone.utc),
restrictions=api_keys_v2.Restrictions(
api_targets=[
api_keys_v2.ApiTarget(service="generativelanguage.googleapis.com")
]
),
)
try: try:
api_keys_client = api_keys_v2.ApiKeysClient(credentials=credentials) api_keys_client = api_keys_v2.ApiKeysClient(credentials=credentials)
api_target = api_keys_v2.ApiTarget(service="generativelanguage.googleapis.com") api_target = api_keys_v2.ApiTarget(service="generativelanguage.googleapis.com")
@@ -313,19 +354,19 @@ def create_api_key(project_id, credentials):
parent=f"projects/{project_id}/locations/global", parent=f"projects/{project_id}/locations/global",
key=key, key=key,
) )
print(" Creating API key...") logging.info(" Creating API key...")
operation = api_keys_client.create_key(request=request) operation = api_keys_client.create_key(request=request)
result = operation.result() result = operation.result()
print(f" Successfully created restricted API key for project {project_id}") logging.info(f" Successfully created restricted API key for project {project_id}")
return result return result
except google_exceptions.PermissionDenied: except google_exceptions.PermissionDenied:
print(f" Permission denied to create API key for project {project_id}. Skipping.") logging.warning(f" Permission denied to create API key for project {project_id}. Skipping.")
return None return None
except google_exceptions.GoogleAPICallError as err: except google_exceptions.GoogleAPICallError as err:
print(f" Error creating API key for project {project_id}: {err}") logging.error(f" Error creating API key for project {project_id}: {err}")
return None return None
def delete_api_keys(project_id, credentials): def delete_api_keys(project_id, credentials, dry_run=False):
"""Deletes all API keys with the display name 'Gemini API Key' and returns their UIDs.""" """Deletes all API keys with the display name 'Gemini API Key' and returns their UIDs."""
deleted_keys_uids = [] deleted_keys_uids = []
try: try:
@@ -336,24 +377,28 @@ def delete_api_keys(project_id, credentials):
keys_to_delete = [key for key in keys if key.display_name == "Gemini API Key"] keys_to_delete = [key for key in keys if key.display_name == "Gemini API Key"]
if not keys_to_delete: if not keys_to_delete:
print(f" No 'Gemini API Key' found to delete.") logging.info(f" No 'Gemini API Key' found to delete.")
return [] return []
print(f" Found {len(keys_to_delete)} key(s) with display name 'Gemini API Key'. Deleting...") logging.info(f" Found {len(keys_to_delete)} key(s) with display name 'Gemini API Key'. Deleting...")
for key in keys_to_delete: for key in keys_to_delete:
if dry_run:
logging.info(f" [DRY RUN] Would delete key: {key.uid}")
deleted_keys_uids.append(key.uid)
continue
try: try:
request = api_keys_v2.DeleteKeyRequest(name=key.name) request = api_keys_v2.DeleteKeyRequest(name=key.name)
operation = api_keys_client.delete_key(request=request) operation = api_keys_client.delete_key(request=request)
operation.result() operation.result()
print(f" Successfully deleted key: {key.uid}") logging.info(f" Successfully deleted key: {key.uid}")
deleted_keys_uids.append(key.uid) deleted_keys_uids.append(key.uid)
except google_exceptions.GoogleAPICallError as err: except google_exceptions.GoogleAPICallError as err:
print(f" Error deleting key {key.uid}: {err}") logging.error(f" Error deleting key {key.uid}: {err}")
return deleted_keys_uids return deleted_keys_uids
except google_exceptions.PermissionDenied: except google_exceptions.PermissionDenied:
print(f" Permission denied to list or delete API keys for project {project_id}. Skipping.") logging.warning(f" Permission denied to list or delete API keys for project {project_id}. Skipping.")
except google_exceptions.GoogleAPICallError as err: except google_exceptions.GoogleAPICallError as err:
print(f" An API error occurred while deleting keys for project {project_id}: {err}") logging.error(f" An API error occurred while deleting keys for project {project_id}: {err}")
return [] return []
if __name__ == "__main__": if __name__ == "__main__":

2
requirements.txt
View File

@@ -2,4 +2,4 @@ google-auth-oauthlib>=1.2.2
google-cloud-api-keys>=0.5.17 google-cloud-api-keys>=0.5.17
google-cloud-resource-manager>=1.14.2 google-cloud-resource-manager>=1.14.2
google-cloud-service-usage>=1.13.1 google-cloud-service-usage>=1.13.1
jsonschema>=4.22.0 jsonschema>=4.25.1

4
schemas/v1/api_keys_database.schema.json
View File

@@ -11,6 +11,10 @@
"type": "string", "type": "string",
"format": "date-time" "format": "date-time"
}, },
"last_modified_utc": {
"type": "string",
"format": "date-time"
},
"accounts": { "accounts": {
"type": "array", "type": "array",
"items": { "items": {