|
1 | | -import configparser |
2 | 1 | import logging |
3 | 2 | import os |
4 | | -from pathlib import Path |
5 | | -from typing import Dict, Optional # Added Dict |
| 3 | +from typing import Optional |
6 | 4 |
|
7 | 5 | import requests |
8 | 6 |
|
9 | 7 | logger = logging.getLogger(__name__) |
10 | 8 |
|
11 | | -# Default locations (used for tests and as fallback). Actual resolution is dynamic via _get_auth_ini_file(). |
12 | | -FIREWORKS_CONFIG_DIR = Path.home() / ".fireworks" |
13 | | -AUTH_INI_FILE = FIREWORKS_CONFIG_DIR / "auth.ini" |
14 | | - |
15 | | - |
16 | | -def _get_profile_base_dir() -> Path: |
17 | | - """ |
18 | | - Resolve the Fireworks configuration base directory following firectl behavior: |
19 | | - - Default: ~/.fireworks |
20 | | - - If FIREWORKS_PROFILE is set and non-empty: ~/.fireworks/profiles/<profile> |
21 | | - """ |
22 | | - profile_name = os.environ.get("FIREWORKS_PROFILE", "").strip() |
23 | | - base_dir = Path.home() / ".fireworks" |
24 | | - if profile_name: |
25 | | - base_dir = base_dir / "profiles" / profile_name |
26 | | - return base_dir |
27 | | - |
28 | | - |
29 | | -def _get_auth_ini_file() -> Path: |
30 | | - """ |
31 | | - Determine the auth.ini file path. |
32 | | - Priority: |
33 | | - 1) FIREWORKS_AUTH_FILE env var when set |
34 | | - 2) ~/.fireworks[/profiles/<profile>]/auth.ini (profile driven) |
35 | | - """ |
36 | | - auth_file_env = os.environ.get("FIREWORKS_AUTH_FILE") |
37 | | - if auth_file_env: |
38 | | - return Path(auth_file_env) |
39 | | - return _get_profile_base_dir() / "auth.ini" |
40 | | - |
41 | | - |
42 | | -def _is_profile_active() -> bool: |
43 | | - """ |
44 | | - Returns True if a specific profile or explicit auth file is active. |
45 | | - In this case, profile-based credentials should take precedence over env vars. |
46 | | - """ |
47 | | - if os.environ.get("FIREWORKS_AUTH_FILE"): |
48 | | - return True |
49 | | - prof = os.environ.get("FIREWORKS_PROFILE", "").strip() |
50 | | - return bool(prof) |
51 | | - |
52 | | - |
53 | | -def _parse_simple_auth_file(file_path: Path) -> Dict[str, str]: |
54 | | - """ |
55 | | - Parses an auth file with simple key=value lines. |
56 | | - Handles comments starting with # or ;. |
57 | | - Strips whitespace and basic quotes from values. |
58 | | - """ |
59 | | - creds = {} |
60 | | - if not file_path.exists(): |
61 | | - return creds |
62 | | - try: |
63 | | - with open(file_path, "r", encoding="utf-8") as f: |
64 | | - for line in f: |
65 | | - line = line.strip() |
66 | | - if not line or line.startswith("#") or line.startswith(";"): |
67 | | - continue |
68 | | - if "=" in line: |
69 | | - key, value = line.split("=", 1) |
70 | | - key = key.strip() |
71 | | - value = value.strip() |
72 | | - # Remove surrounding quotes if present |
73 | | - if value and ( |
74 | | - (value.startswith('"') and value.endswith('"')) |
75 | | - or (value.startswith("'") and value.endswith("'")) |
76 | | - ): |
77 | | - value = value[1:-1] |
78 | | - |
79 | | - if key in ["api_key", "account_id"] and value: |
80 | | - creds[key] = value |
81 | | - except Exception as e: |
82 | | - logger.warning("Error during simple parsing of %s: %s", str(file_path), e) |
83 | | - return creds |
84 | | - |
85 | | - |
86 | | -def _get_credential_from_config_file(key_name: str) -> Optional[str]: |
87 | | - """ |
88 | | - Helper to get a specific credential (api_key or account_id) from auth.ini. |
89 | | - Tries simple parsing first, then configparser. |
90 | | - """ |
91 | | - auth_ini_path = _get_auth_ini_file() |
92 | | - if not auth_ini_path.exists(): |
93 | | - return None |
94 | | - |
95 | | - # 1. Try simple key-value parsing first |
96 | | - simple_creds = _parse_simple_auth_file(auth_ini_path) |
97 | | - if key_name in simple_creds: |
98 | | - logger.debug("Using %s from simple key-value parsing of %s.", key_name, str(auth_ini_path)) |
99 | | - return simple_creds[key_name] |
100 | | - |
101 | | - # 2. Fallback to configparser if not found via simple parsing or if simple parsing failed |
102 | | - # This path will also generate the "no section headers" warning if applicable, |
103 | | - # but only if simple parsing didn't yield the key. |
104 | | - try: |
105 | | - config = configparser.ConfigParser() |
106 | | - config.read(auth_ini_path) |
107 | | - |
108 | | - # Try [fireworks] section |
109 | | - if "fireworks" in config and config.has_option("fireworks", key_name): |
110 | | - value_from_file = config.get("fireworks", key_name) |
111 | | - if value_from_file: |
112 | | - logger.debug("Using %s from [fireworks] section in %s.", key_name, str(auth_ini_path)) |
113 | | - return value_from_file |
114 | | - |
115 | | - # Try default section (configparser might place items without section header here) |
116 | | - if config.has_option(config.default_section, key_name): |
117 | | - value_from_default = config.get(config.default_section, key_name) |
118 | | - if value_from_default: |
119 | | - logger.debug( |
120 | | - "Using %s from default section [%s] in %s.", |
121 | | - key_name, |
122 | | - config.default_section, |
123 | | - str(auth_ini_path), |
124 | | - ) |
125 | | - return value_from_default |
126 | | - |
127 | | - except configparser.MissingSectionHeaderError: |
128 | | - # This error implies the file is purely key-value, which simple parsing should have handled. |
129 | | - # If simple parsing failed to get the key, then it's likely not there or malformed. |
130 | | - logger.debug("%s has no section headers, and simple parsing did not find %s.", str(auth_ini_path), key_name) |
131 | | - except configparser.Error as e_config: |
132 | | - logger.warning("Configparser error reading %s for %s: %s", str(auth_ini_path), key_name, e_config) |
133 | | - except Exception as e_general: |
134 | | - logger.warning("Unexpected error reading %s for %s: %s", str(auth_ini_path), key_name, e_general) |
135 | | - |
136 | | - return None |
137 | | - |
138 | | - |
139 | | -def _get_credentials_from_config_file() -> Dict[str, Optional[str]]: |
140 | | - """ |
141 | | - Retrieve both api_key and account_id from auth.ini with a single read/parse. |
142 | | - Tries simple parsing first for both keys, then falls back to configparser for any missing ones. |
143 | | - Returns a dict with up to two keys: 'api_key' and 'account_id'. |
144 | | - """ |
145 | | - results: Dict[str, Optional[str]] = {} |
146 | | - auth_ini_path = _get_auth_ini_file() |
147 | | - if not auth_ini_path.exists(): |
148 | | - return results |
149 | | - |
150 | | - # 1) Simple key=value parsing |
151 | | - try: |
152 | | - simple_creds = _parse_simple_auth_file(auth_ini_path) |
153 | | - if "api_key" in simple_creds and simple_creds["api_key"]: |
154 | | - results["api_key"] = simple_creds["api_key"] |
155 | | - if "account_id" in simple_creds and simple_creds["account_id"]: |
156 | | - results["account_id"] = simple_creds["account_id"] |
157 | | - if "api_key" in results and "account_id" in results: |
158 | | - return results |
159 | | - except Exception as e: |
160 | | - logger.warning("Error during simple parsing of %s: %s", str(auth_ini_path), e) |
161 | | - |
162 | | - # 2) ConfigParser for any missing keys |
163 | | - try: |
164 | | - config = configparser.ConfigParser() |
165 | | - config.read(auth_ini_path) |
166 | | - for key_name in ("api_key", "account_id"): |
167 | | - if key_name in results and results[key_name]: |
168 | | - continue |
169 | | - if "fireworks" in config and config.has_option("fireworks", key_name): |
170 | | - value_from_file = config.get("fireworks", key_name) |
171 | | - if value_from_file: |
172 | | - results[key_name] = value_from_file |
173 | | - continue |
174 | | - if config.has_option(config.default_section, key_name): |
175 | | - value_from_default = config.get(config.default_section, key_name) |
176 | | - if value_from_default: |
177 | | - results[key_name] = value_from_default |
178 | | - except configparser.MissingSectionHeaderError: |
179 | | - # Purely key=value file without section headers; simple parsing should have handled it already. |
180 | | - logger.debug("%s has no section headers; falling back to simple parsing results.", str(auth_ini_path)) |
181 | | - except configparser.Error as e_config: |
182 | | - logger.warning("Configparser error reading %s: %s", str(auth_ini_path), e_config) |
183 | | - except Exception as e_general: |
184 | | - logger.warning("Unexpected error reading %s: %s", str(auth_ini_path), e_general) |
185 | | - |
186 | | - return results |
187 | | - |
188 | 9 |
|
189 | 10 | def get_fireworks_api_key() -> Optional[str]: |
190 | 11 | """ |
191 | 12 | Retrieves the Fireworks API key. |
192 | 13 |
|
193 | | - The key is sourced in the following order: |
194 | | - 1. FIREWORKS_API_KEY environment variable. |
195 | | - 2. 'api_key' from the [fireworks] section of ~/.fireworks/auth.ini. |
196 | | -
|
197 | 14 | Returns: |
198 | 15 | The API key if found, otherwise None. |
199 | 16 | """ |
200 | | - # If a profile is active, prefer profile file first, then env |
201 | | - if _is_profile_active(): |
202 | | - api_key_from_file = _get_credential_from_config_file("api_key") |
203 | | - if api_key_from_file: |
204 | | - return api_key_from_file |
205 | | - api_key = os.environ.get("FIREWORKS_API_KEY") |
206 | | - if api_key: |
207 | | - logger.debug("Using FIREWORKS_API_KEY from environment variable (profile active but file missing).") |
208 | | - return api_key |
209 | | - else: |
210 | | - # Default behavior: env overrides file |
211 | | - api_key = os.environ.get("FIREWORKS_API_KEY") |
212 | | - if api_key: |
213 | | - logger.debug("Using FIREWORKS_API_KEY from environment variable.") |
214 | | - return api_key |
215 | | - api_key_from_file = _get_credential_from_config_file("api_key") |
216 | | - if api_key_from_file: |
217 | | - return api_key_from_file |
218 | | - |
219 | | - logger.debug("Fireworks API key not found in environment variables or auth.ini.") |
| 17 | + api_key = os.environ.get("FIREWORKS_API_KEY") |
| 18 | + if api_key and api_key.strip(): |
| 19 | + logger.debug("Using FIREWORKS_API_KEY from environment variable.") |
| 20 | + return api_key.strip() |
| 21 | + logger.debug("Fireworks API key not found in environment variables.") |
220 | 22 | return None |
221 | 23 |
|
222 | 24 |
|
223 | 25 | def get_fireworks_account_id() -> Optional[str]: |
224 | 26 | """ |
225 | 27 | Retrieves the Fireworks Account ID. |
226 | 28 |
|
227 | | - The Account ID is sourced in the following order: |
228 | | - 1. FIREWORKS_ACCOUNT_ID environment variable. |
229 | | - 2. 'account_id' from the [fireworks] section of ~/.fireworks/auth.ini. |
230 | | - 3. If an API key is available (env or auth.ini), resolve via verifyApiKey. |
231 | | -
|
232 | 29 | Returns: |
233 | 30 | The Account ID if found, otherwise None. |
234 | 31 | """ |
235 | | - # If a profile is active, prefer profile file first, then env |
236 | | - if _is_profile_active(): |
237 | | - creds = _get_credentials_from_config_file() |
238 | | - account_id_from_file = creds.get("account_id") |
239 | | - if account_id_from_file: |
240 | | - return account_id_from_file |
241 | | - account_id = os.environ.get("FIREWORKS_ACCOUNT_ID") |
242 | | - if account_id: |
243 | | - logger.debug("Using FIREWORKS_ACCOUNT_ID from environment variable (profile active but file missing).") |
244 | | - return account_id |
245 | | - else: |
246 | | - # Default behavior: env overrides file |
247 | | - account_id = os.environ.get("FIREWORKS_ACCOUNT_ID") |
248 | | - if account_id: |
249 | | - logger.debug("Using FIREWORKS_ACCOUNT_ID from environment variable.") |
250 | | - return account_id |
251 | | - creds = _get_credentials_from_config_file() |
252 | | - account_id_from_file = creds.get("account_id") |
253 | | - if account_id_from_file: |
254 | | - return account_id_from_file |
255 | | - |
256 | | - # 3) Fallback: if API key is present, attempt to resolve via verifyApiKey (env or auth.ini) |
| 32 | + # Account id is derived from the API key (single source of truth). |
257 | 33 | try: |
258 | | - # Intentionally use get_fireworks_api_key to centralize precedence (env vs file) |
259 | 34 | api_key_for_verify = get_fireworks_api_key() |
260 | 35 | if api_key_for_verify: |
261 | 36 | resolved = verify_api_key_and_get_account_id(api_key=api_key_for_verify, api_base=get_fireworks_api_base()) |
262 | 37 | if resolved: |
263 | | - logger.debug("Using FIREWORKS_ACCOUNT_ID resolved via verifyApiKey: %s", resolved) |
| 38 | + logger.debug("Resolved account id via verifyApiKey: %s", resolved) |
264 | 39 | return resolved |
265 | 40 | except Exception as e: |
266 | | - logger.debug("Failed to resolve FIREWORKS_ACCOUNT_ID via verifyApiKey: %s", e) |
| 41 | + logger.debug("Failed to resolve account id via verifyApiKey: %s", e) |
267 | 42 |
|
268 | | - logger.debug("Fireworks Account ID not found in environment variables, auth.ini, or via verifyApiKey.") |
| 43 | + logger.debug("Fireworks Account ID not found via verifyApiKey.") |
269 | 44 | return None |
270 | 45 |
|
271 | 46 |
|
@@ -323,7 +98,7 @@ def verify_api_key_and_get_account_id( |
323 | 98 | # Header keys could vary in case; requests provides case-insensitive dict |
324 | 99 | account_id = resp.headers.get("x-fireworks-account-id") or resp.headers.get("X-Fireworks-Account-Id") |
325 | 100 | if account_id and account_id.strip(): |
326 | | - logger.debug("Resolved FIREWORKS_ACCOUNT_ID via verifyApiKey: %s", account_id) |
| 101 | + logger.debug("Resolved account id via verifyApiKey: %s", account_id) |
327 | 102 | return account_id.strip() |
328 | 103 | return None |
329 | 104 | except Exception as e: |
|
0 commit comments