mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
121 lines
4.7 KiB
Python
121 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Ban hardcoded commercial-partition AWS hosts and ARN prefixes under `litellm/`.
|
|
|
|
An endpoint or ARN built with a literal `amazonaws.com` or `arn:aws:` works in every
|
|
commercial region and breaks only for GovCloud (`us-gov-*`, `arn:aws-us-gov:`) and
|
|
China (`amazonaws.com.cn`, `arn:aws-cn:`) deployments, so the failure never shows up
|
|
in CI or on a developer laptop. `litellm/litellm_core_utils/aws_partition.py` derives
|
|
both from the region and is the only place those literals belong. Build hosts with
|
|
`get_aws_dns_suffix(region)` and ARNs with `get_aws_arn_prefix(region)`.
|
|
|
|
Every string constant in every `litellm/**/*.py` file is scanned, including the
|
|
literal parts of f-strings and the strings inside `.format()` calls and
|
|
concatenations. Docstrings and comments are not, since they never reach a request.
|
|
`amazonaws.com.cn` passes because it is already the China partition.
|
|
|
|
`ALLOWED` holds the (file, token) pairs that are text rather than a request target:
|
|
a hosted logo, an IAM service principal, and hostnames quoted as examples inside
|
|
error messages and field descriptions. An entry only covers that exact token in that
|
|
exact file, so a second literal in an allowed file is still caught, and an entry
|
|
whose token is gone fails the check so the set only shrinks.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Final, NamedTuple
|
|
|
|
REPO_ROOT: Final = Path(__file__).resolve().parents[2]
|
|
SCAN_ROOT: Final = REPO_ROOT / "litellm"
|
|
PARTITION_HELPER: Final = "litellm/litellm_core_utils/aws_partition.py"
|
|
|
|
COMMERCIAL_TOKEN: Final = re.compile(r"[A-Za-z0-9.-]*amazonaws\.com(?!\.cn)|arn:aws:[A-Za-z0-9:/_.*-]*")
|
|
|
|
|
|
class Allowance(NamedTuple):
|
|
file: str
|
|
token: str
|
|
|
|
|
|
ALLOWED: Final = frozenset(
|
|
{
|
|
Allowance("litellm/integrations/email_alerting.py", "litellm-listing.s3.amazonaws.com"),
|
|
Allowance("litellm/types/integrations/slack_alerting.py", "litellm-listing.s3.amazonaws.com"),
|
|
Allowance("litellm/rag/ingestion/bedrock_ingestion.py", "bedrock.amazonaws.com"),
|
|
Allowance(
|
|
"litellm/llms/bedrock/chat/agentcore/transformation.py",
|
|
"arn:aws:bedrock-agentcore:region:account:runtime/runtime_id",
|
|
),
|
|
Allowance("litellm/llms/bedrock/search/transformation.py", ".amazonaws.com"),
|
|
Allowance(
|
|
"litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py",
|
|
"bucket.s3.amazonaws.com",
|
|
),
|
|
Allowance("litellm/types/proxy/claude_code_endpoints.py", "bucket.s3.amazonaws.com"),
|
|
}
|
|
)
|
|
|
|
|
|
class Hit(NamedTuple):
|
|
file: str
|
|
line: int
|
|
token: str
|
|
|
|
|
|
def _docstring_ids(tree: ast.Module) -> frozenset[int]:
|
|
return frozenset(
|
|
id(statement.value)
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef))
|
|
for statement in node.body
|
|
if isinstance(statement, ast.Expr)
|
|
and isinstance(statement.value, ast.Constant)
|
|
and isinstance(statement.value.value, str)
|
|
)
|
|
|
|
|
|
def _hits_in_file(path: Path) -> tuple[Hit, ...]:
|
|
tree: Final = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
|
docstrings: Final = _docstring_ids(tree)
|
|
relative: Final = path.relative_to(REPO_ROOT).as_posix()
|
|
return tuple(
|
|
Hit(relative, node.lineno, match.group(0))
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstrings
|
|
for match in COMMERCIAL_TOKEN.finditer(node.value)
|
|
)
|
|
|
|
|
|
def find_hits(scan_root: Path) -> tuple[Hit, ...]:
|
|
return tuple(
|
|
hit
|
|
for path in sorted(scan_root.rglob("*.py"))
|
|
if path.relative_to(REPO_ROOT).as_posix() != PARTITION_HELPER
|
|
for hit in _hits_in_file(path)
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
hits: Final = find_hits(SCAN_ROOT)
|
|
seen: Final = frozenset(Allowance(hit.file, hit.token) for hit in hits)
|
|
violations: Final = tuple(hit for hit in hits if Allowance(hit.file, hit.token) not in ALLOWED)
|
|
stale: Final = ALLOWED - seen
|
|
for hit in violations:
|
|
print(f"{hit.file}:{hit.line}: hardcoded commercial AWS partition literal {hit.token!r}")
|
|
for allowance in sorted(stale):
|
|
print(f"{allowance.file}: ALLOWED entry {allowance.token!r} no longer matches anything, remove it")
|
|
if violations or stale:
|
|
print(
|
|
"\nBuild AWS hosts with get_aws_dns_suffix(region) and ARNs with get_aws_arn_prefix(region) "
|
|
"from litellm/litellm_core_utils/aws_partition.py so GovCloud and China regions resolve."
|
|
)
|
|
return 1
|
|
print(f"No hardcoded commercial AWS partition literals outside {PARTITION_HELPER}.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|