Spaces:
Running
Running
File size: 12,410 Bytes
7cc3183 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 |
import os
import glob
import random
import json
from typing import List, Dict, Any
from google.auth.transport.requests import Request as AuthRequest
from google.oauth2 import service_account
import config as app_config # Changed from relative
# Helper function to parse multiple JSONs from a string
def parse_multiple_json_credentials(json_str: str) -> List[Dict[str, Any]]:
"""
Parse multiple JSON objects from a string separated by commas.
Format expected: {json_object1},{json_object2},...
Returns a list of parsed JSON objects.
"""
credentials_list = []
nesting_level = 0
current_object_start = -1
str_length = len(json_str)
for i, char in enumerate(json_str):
if char == '{':
if nesting_level == 0:
current_object_start = i
nesting_level += 1
elif char == '}':
if nesting_level > 0:
nesting_level -= 1
if nesting_level == 0 and current_object_start != -1:
# Found a complete top-level JSON object
json_object_str = json_str[current_object_start : i + 1]
try:
credentials_info = json.loads(json_object_str)
# Basic validation for service account structure
required_fields = ["type", "project_id", "private_key_id", "private_key", "client_email"]
if all(field in credentials_info for field in required_fields):
credentials_list.append(credentials_info)
print(f"DEBUG: Successfully parsed a JSON credential object.")
else:
print(f"WARNING: Parsed JSON object missing required fields: {json_object_str[:100]}...")
except json.JSONDecodeError as e:
print(f"ERROR: Failed to parse JSON object segment: {json_object_str[:100]}... Error: {e}")
current_object_start = -1 # Reset for the next object
else:
# Found a closing brace without a matching open brace in scope, might indicate malformed input
print(f"WARNING: Encountered unexpected '}}' at index {i}. Input might be malformed.")
if nesting_level != 0:
print(f"WARNING: JSON string parsing ended with non-zero nesting level ({nesting_level}). Check for unbalanced braces.")
print(f"DEBUG: Parsed {len(credentials_list)} credential objects from the input string.")
return credentials_list
def _refresh_auth(credentials):
"""Helper function to refresh GCP token."""
if not credentials:
print("ERROR: _refresh_auth called with no credentials.")
return None
try:
# Assuming credentials object has a project_id attribute for logging
project_id_for_log = getattr(credentials, 'project_id', 'Unknown')
print(f"INFO: Attempting to refresh token for project: {project_id_for_log}...")
credentials.refresh(AuthRequest())
print(f"INFO: Token refreshed successfully for project: {project_id_for_log}")
return credentials.token
except Exception as e:
project_id_for_log = getattr(credentials, 'project_id', 'Unknown')
print(f"ERROR: Error refreshing GCP token for project {project_id_for_log}: {e}")
return None
# Credential Manager for handling multiple service accounts
class CredentialManager:
def __init__(self): # default_credentials_dir is now handled by config
# Use CREDENTIALS_DIR from config
self.credentials_dir = app_config.CREDENTIALS_DIR
self.credentials_files = []
self.current_index = 0
self.credentials = None
self.project_id = None
# New: Store credentials loaded directly from JSON objects
self.in_memory_credentials: List[Dict[str, Any]] = []
self.load_credentials_list() # Load file-based credentials initially
def add_credential_from_json(self, credentials_info: Dict[str, Any]) -> bool:
"""
Add a credential from a JSON object to the manager's in-memory list.
Args:
credentials_info: Dict containing service account credentials
Returns:
bool: True if credential was added successfully, False otherwise
"""
try:
# Validate structure again before creating credentials object
required_fields = ["type", "project_id", "private_key_id", "private_key", "client_email"]
if not all(field in credentials_info for field in required_fields):
print(f"WARNING: Skipping JSON credential due to missing required fields.")
return False
credentials = service_account.Credentials.from_service_account_info(
credentials_info,
scopes=['https://www.googleapis.com/auth/cloud-platform']
)
project_id = credentials.project_id
print(f"DEBUG: Successfully created credentials object from JSON for project: {project_id}")
# Store the credentials object and project ID
self.in_memory_credentials.append({
'credentials': credentials,
'project_id': project_id,
'source': 'json_string' # Add source for clarity
})
print(f"INFO: Added credential for project {project_id} from JSON string to Credential Manager.")
return True
except Exception as e:
print(f"ERROR: Failed to create credentials from parsed JSON object: {e}")
return False
def load_credentials_from_json_list(self, json_list: List[Dict[str, Any]]) -> int:
"""
Load multiple credentials from a list of JSON objects into memory.
Args:
json_list: List of dicts containing service account credentials
Returns:
int: Number of credentials successfully loaded
"""
# Avoid duplicates if called multiple times
existing_projects = {cred['project_id'] for cred in self.in_memory_credentials}
success_count = 0
newly_added_projects = set()
for credentials_info in json_list:
project_id = credentials_info.get('project_id')
# Check if this project_id from JSON exists in files OR already added from JSON
is_duplicate_file = any(os.path.basename(f) == f"{project_id}.json" for f in self.credentials_files) # Basic check
is_duplicate_mem = project_id in existing_projects or project_id in newly_added_projects
if project_id and not is_duplicate_file and not is_duplicate_mem:
if self.add_credential_from_json(credentials_info):
success_count += 1
newly_added_projects.add(project_id)
elif project_id:
print(f"DEBUG: Skipping duplicate credential for project {project_id} from JSON list.")
if success_count > 0:
print(f"INFO: Loaded {success_count} new credentials from JSON list into memory.")
return success_count
def load_credentials_list(self):
"""Load the list of available credential files"""
# Look for all .json files in the credentials directory
pattern = os.path.join(self.credentials_dir, "*.json")
self.credentials_files = glob.glob(pattern)
if not self.credentials_files:
# print(f"No credential files found in {self.credentials_dir}")
pass # Don't return False yet, might have in-memory creds
else:
print(f"Found {len(self.credentials_files)} credential files: {[os.path.basename(f) for f in self.credentials_files]}")
# Check total credentials
return self.get_total_credentials() > 0
def refresh_credentials_list(self):
"""Refresh the list of credential files and return if any credentials exist"""
old_file_count = len(self.credentials_files)
self.load_credentials_list() # Reloads file list
new_file_count = len(self.credentials_files)
if old_file_count != new_file_count:
print(f"Credential files updated: {old_file_count} -> {new_file_count}")
# Total credentials = files + in-memory
total_credentials = self.get_total_credentials()
print(f"DEBUG: Refresh check - Total credentials available: {total_credentials}")
return total_credentials > 0
def get_total_credentials(self):
"""Returns the total number of credentials (file + in-memory)."""
return len(self.credentials_files) + len(self.in_memory_credentials)
def get_random_credentials(self):
"""
Get a random credential (file or in-memory) and load it.
Tries each available credential source at most once in a random order.
"""
all_sources = []
# Add file paths (as type 'file')
for file_path in self.credentials_files:
all_sources.append({'type': 'file', 'value': file_path})
# Add in-memory credentials (as type 'memory_object')
# Assuming self.in_memory_credentials stores dicts like {'credentials': cred_obj, 'project_id': pid, 'source': 'json_string'}
for idx, mem_cred_info in enumerate(self.in_memory_credentials):
all_sources.append({'type': 'memory_object', 'value': mem_cred_info, 'original_index': idx})
if not all_sources:
print("WARNING: No credentials available for random selection (no files or in-memory).")
return None, None
random.shuffle(all_sources) # Shuffle to try in a random order
for source_info in all_sources:
source_type = source_info['type']
if source_type == 'file':
file_path = source_info['value']
print(f"DEBUG: Attempting to load credential from file: {os.path.basename(file_path)}")
try:
credentials = service_account.Credentials.from_service_account_file(
file_path,
scopes=['https://www.googleapis.com/auth/cloud-platform']
)
project_id = credentials.project_id
print(f"INFO: Successfully loaded credential from file {os.path.basename(file_path)} for project: {project_id}")
self.credentials = credentials # Cache last successfully loaded
self.project_id = project_id
return credentials, project_id
except Exception as e:
print(f"ERROR: Failed loading credentials file {os.path.basename(file_path)}: {e}. Trying next available source.")
continue # Try next source
elif source_type == 'memory_object':
mem_cred_detail = source_info['value']
# The 'credentials' object is already a service_account.Credentials instance
credentials = mem_cred_detail.get('credentials')
project_id = mem_cred_detail.get('project_id')
if credentials and project_id:
print(f"INFO: Using in-memory credential for project: {project_id} (Source: {mem_cred_detail.get('source', 'unknown')})")
# Here, we might want to ensure the credential object is still valid if it can expire
# For service_account.Credentials from_service_account_info, they typically don't self-refresh
# in the same way as ADC, but are long-lived based on the private key.
# If validation/refresh were needed, it would be complex here.
# For now, assume it's usable if present.
self.credentials = credentials # Cache last successfully loaded/used
self.project_id = project_id
return credentials, project_id
else:
print(f"WARNING: In-memory credential entry missing 'credentials' or 'project_id' at original index {source_info.get('original_index', 'N/A')}. Skipping.")
continue # Try next source
print("WARNING: All available credential sources failed to load.")
return None, None |