sana0721/vertex
0
1import os2import glob3import random4import json5from typing import List, Dict, Any6from google.auth.transport.requests import Request as AuthRequest7from google.oauth2 import service_account8import config as app_config # Changed from relative9 10# Helper function to parse multiple JSONs from a string11def parse_multiple_json_credentials(json_str: str) -> List[Dict[str, Any]]:12 """13 Parse multiple JSON objects from a string separated by commas.14 Format expected: {json_object1},{json_object2},...15 Returns a list of parsed JSON objects.16 """17 credentials_list = []18 nesting_level = 019 current_object_start = -120 str_length = len(json_str)21 22 for i, char in enumerate(json_str):23 if char == '{':24 if nesting_level == 0:25 current_object_start = i26 nesting_level += 127 elif char == '}':28 if nesting_level > 0:29 nesting_level -= 130 if nesting_level == 0 and current_object_start != -1:31 # Found a complete top-level JSON object32 json_object_str = json_str[current_object_start : i + 1]33 try:34 credentials_info = json.loads(json_object_str)35 # Basic validation for service account structure36 required_fields = ["type", "project_id", "private_key_id", "private_key", "client_email"]37 if all(field in credentials_info for field in required_fields):38 credentials_list.append(credentials_info)39 print(f"DEBUG: Successfully parsed a JSON credential object.")40 else:41 print(f"WARNING: Parsed JSON object missing required fields: {json_object_str[:100]}...")42 except json.JSONDecodeError as e:43 print(f"ERROR: Failed to parse JSON object segment: {json_object_str[:100]}... Error: {e}")44 current_object_start = -1 # Reset for the next object45 else:46 # Found a closing brace without a matching open brace in scope, might indicate malformed input47 print(f"WARNING: Encountered unexpected '}}' at index {i}. Input might be malformed.")48 49 50 if nesting_level != 0:51 print(f"WARNING: JSON string parsing ended with non-zero nesting level ({nesting_level}). Check for unbalanced braces.")52 53 print(f"DEBUG: Parsed {len(credentials_list)} credential objects from the input string.")54 return credentials_list55def _refresh_auth(credentials):56 """Helper function to refresh GCP token."""57 if not credentials:58 print("ERROR: _refresh_auth called with no credentials.")59 return None60 try:61 # Assuming credentials object has a project_id attribute for logging62 project_id_for_log = getattr(credentials, 'project_id', 'Unknown')63 print(f"INFO: Attempting to refresh token for project: {project_id_for_log}...")64 credentials.refresh(AuthRequest())65 print(f"INFO: Token refreshed successfully for project: {project_id_for_log}")66 return credentials.token67 except Exception as e:68 project_id_for_log = getattr(credentials, 'project_id', 'Unknown')69 print(f"ERROR: Error refreshing GCP token for project {project_id_for_log}: {e}")70 return None71 72 73# Credential Manager for handling multiple service accounts74class CredentialManager:75 def __init__(self): # default_credentials_dir is now handled by config76 # Use CREDENTIALS_DIR from config77 self.credentials_dir = app_config.CREDENTIALS_DIR78 self.credentials_files = []79 self.current_index = 080 self.credentials = None81 self.project_id = None82 # New: Store credentials loaded directly from JSON objects83 self.in_memory_credentials: List[Dict[str, Any]] = []84 # Round-robin index for tracking position85 self.round_robin_index = 086 self.load_credentials_list() # Load file-based credentials initially87 88 def add_credential_from_json(self, credentials_info: Dict[str, Any]) -> bool:89 """90 Add a credential from a JSON object to the manager's in-memory list.91 92 Args:93 credentials_info: Dict containing service account credentials94 95 Returns:96 bool: True if credential was added successfully, False otherwise97 """98 try:99 # Validate structure again before creating credentials object100 required_fields = ["type", "project_id", "private_key_id", "private_key", "client_email"]101 if not all(field in credentials_info for field in required_fields):102 print(f"WARNING: Skipping JSON credential due to missing required fields.")103 return False104 105 credentials = service_account.Credentials.from_service_account_info(106 credentials_info,107 scopes=['https://www.googleapis.com/auth/cloud-platform']108 )109 project_id = credentials.project_id110 print(f"DEBUG: Successfully created credentials object from JSON for project: {project_id}")111 112 # Store the credentials object and project ID113 self.in_memory_credentials.append({114 'credentials': credentials,115 'project_id': project_id,116 'source': 'json_string' # Add source for clarity117 })118 print(f"INFO: Added credential for project {project_id} from JSON string to Credential Manager.")119 return True120 except Exception as e:121 print(f"ERROR: Failed to create credentials from parsed JSON object: {e}")122 return False123 124 def load_credentials_from_json_list(self, json_list: List[Dict[str, Any]]) -> int:125 """126 Load multiple credentials from a list of JSON objects into memory.127 128 Args:129 json_list: List of dicts containing service account credentials130 131 Returns:132 int: Number of credentials successfully loaded133 """134 # Avoid duplicates if called multiple times135 existing_projects = {cred['project_id'] for cred in self.in_memory_credentials}136 success_count = 0137 newly_added_projects = set()138 139 for credentials_info in json_list:140 project_id = credentials_info.get('project_id')141 # Check if this project_id from JSON exists in files OR already added from JSON142 is_duplicate_file = any(os.path.basename(f) == f"{project_id}.json" for f in self.credentials_files) # Basic check143 is_duplicate_mem = project_id in existing_projects or project_id in newly_added_projects144 145 if project_id and not is_duplicate_file and not is_duplicate_mem:146 if self.add_credential_from_json(credentials_info):147 success_count += 1148 newly_added_projects.add(project_id)149 elif project_id:150 print(f"DEBUG: Skipping duplicate credential for project {project_id} from JSON list.")151 152 153 if success_count > 0:154 print(f"INFO: Loaded {success_count} new credentials from JSON list into memory.")155 return success_count156 157 def load_credentials_list(self):158 """Load the list of available credential files"""159 # Look for all .json files in the credentials directory160 pattern = os.path.join(self.credentials_dir, "*.json")161 self.credentials_files = glob.glob(pattern)162 163 if not self.credentials_files:164 # print(f"No credential files found in {self.credentials_dir}")165 pass # Don't return False yet, might have in-memory creds166 else:167 print(f"Found {len(self.credentials_files)} credential files: {[os.path.basename(f) for f in self.credentials_files]}")168 169 # Check total credentials170 return self.get_total_credentials() > 0171 172 def refresh_credentials_list(self):173 """Refresh the list of credential files and return if any credentials exist"""174 old_file_count = len(self.credentials_files)175 self.load_credentials_list() # Reloads file list176 new_file_count = len(self.credentials_files)177 178 if old_file_count != new_file_count:179 print(f"Credential files updated: {old_file_count} -> {new_file_count}")180 181 # Total credentials = files + in-memory182 total_credentials = self.get_total_credentials()183 print(f"DEBUG: Refresh check - Total credentials available: {total_credentials}")184 return total_credentials > 0185 186 def get_total_credentials(self):187 """Returns the total number of credentials (file + in-memory)."""188 return len(self.credentials_files) + len(self.in_memory_credentials)189 190 191 def _get_all_credential_sources(self):192 """193 Get all available credential sources (files and in-memory).194 Returns a list of dicts with 'type' and 'value' keys.195 """196 all_sources = []197 198 # Add file paths (as type 'file')199 for file_path in self.credentials_files:200 all_sources.append({'type': 'file', 'value': file_path})201 202 # Add in-memory credentials (as type 'memory_object')203 for idx, mem_cred_info in enumerate(self.in_memory_credentials):204 all_sources.append({'type': 'memory_object', 'value': mem_cred_info, 'original_index': idx})205 206 return all_sources207 208 def _load_credential_from_source(self, source_info):209 """210 Load a credential from a given source.211 Returns (credentials, project_id) tuple or (None, None) on failure.212 """213 source_type = source_info['type']214 215 if source_type == 'file':216 file_path = source_info['value']217 print(f"DEBUG: Attempting to load credential from file: {os.path.basename(file_path)}")218 try:219 credentials = service_account.Credentials.from_service_account_file(220 file_path,221 scopes=['https://www.googleapis.com/auth/cloud-platform']222 )223 project_id = credentials.project_id224 print(f"INFO: Successfully loaded credential from file {os.path.basename(file_path)} for project: {project_id}")225 self.credentials = credentials # Cache last successfully loaded226 self.project_id = project_id227 return credentials, project_id228 except Exception as e:229 print(f"ERROR: Failed loading credentials file {os.path.basename(file_path)}: {e}")230 return None, None231 232 elif source_type == 'memory_object':233 mem_cred_detail = source_info['value']234 credentials = mem_cred_detail.get('credentials')235 project_id = mem_cred_detail.get('project_id')236 237 if credentials and project_id:238 print(f"INFO: Using in-memory credential for project: {project_id} (Source: {mem_cred_detail.get('source', 'unknown')})")239 self.credentials = credentials # Cache last successfully loaded/used240 self.project_id = project_id241 return credentials, project_id242 else:243 print(f"WARNING: In-memory credential entry missing 'credentials' or 'project_id' at original index {source_info.get('original_index', 'N/A')}.")244 return None, None245 246 return None, None247 248 def get_random_credentials(self):249 """250 Get a random credential from available sources.251 Tries each available credential source at most once in random order.252 Returns (credentials, project_id) tuple or (None, None) if all fail.253 """254 all_sources = self._get_all_credential_sources()255 256 if not all_sources:257 print("WARNING: No credentials available for selection (no files or in-memory).")258 return None, None259 260 print(f"DEBUG: Using random credential selection strategy.")261 sources_to_try = all_sources.copy()262 random.shuffle(sources_to_try) # Shuffle to try in a random order263 264 for source_info in sources_to_try:265 credentials, project_id = self._load_credential_from_source(source_info)266 if credentials and project_id:267 return credentials, project_id268 269 print("WARNING: All available credential sources failed to load.")270 return None, None271 272 def get_roundrobin_credentials(self):273 """274 Get a credential using round-robin selection.275 Tries credentials in order, cycling through all available sources.276 Returns (credentials, project_id) tuple or (None, None) if all fail.277 """278 all_sources = self._get_all_credential_sources()279 280 if not all_sources:281 print("WARNING: No credentials available for selection (no files or in-memory).")282 return None, None283 284 print(f"DEBUG: Using round-robin credential selection strategy.")285 286 # Ensure round_robin_index is within bounds287 if self.round_robin_index >= len(all_sources):288 self.round_robin_index = 0289 290 # Create ordered list starting from round_robin_index291 ordered_sources = all_sources[self.round_robin_index:] + all_sources[:self.round_robin_index]292 293 # Move to next index for next call294 self.round_robin_index = (self.round_robin_index + 1) % len(all_sources)295 296 # Try credentials in round-robin order297 for source_info in ordered_sources:298 credentials, project_id = self._load_credential_from_source(source_info)299 if credentials and project_id:300 return credentials, project_id301 302 print("WARNING: All available credential sources failed to load.")303 return None, None304 305 def get_credentials(self):306 """307 Get credentials based on the configured selection strategy.308 Checks ROUNDROBIN config and calls the appropriate method.309 Returns (credentials, project_id) tuple or (None, None) if all fail.310 """311 if app_config.ROUNDROBIN:312 return self.get_roundrobin_credentials()313 else:314 return self.get_random_credentials()