simpleaudit 0.3.1__tar.gz → 0.3.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {simpleaudit-0.3.1/simpleaudit.egg-info → simpleaudit-0.3.3}/PKG-INFO +1 -1
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/pyproject.toml +1 -1
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/__init__.py +2 -1
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/model_auditor.py +55 -1
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/model.py +14 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/__init__.py +6 -0
- simpleaudit-0.3.3/simpleaudit/tracing/auth.py +332 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/otlp.py +43 -13
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/provider.py +11 -2
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/selection.py +29 -3
- {simpleaudit-0.3.1 → simpleaudit-0.3.3/simpleaudit.egg-info}/PKG-INFO +1 -1
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/SOURCES.txt +1 -0
- simpleaudit-0.3.3/tests/test_otlp_auth.py +401 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_targets.py +69 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_tracing.py +22 -1
- simpleaudit-0.3.1/simpleaudit/tracing/auth.py +0 -175
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/LICENSE +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/README.md +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/setup.cfg +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/auditor.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/checklist.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/cli.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/context_marks.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/cross_judge.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/experiment.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/__init__.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/abstention.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/binary_abstention.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/checklist.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/compose.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/default.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/factuality.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/harm.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/helpfulness.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/helsedir_sexhealth_no.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/helsedir_sexhealth_no_rag.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/safety.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/perturbations.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/reframing.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/repeated_results.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/results.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/__init__.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/arbeidstilsynet_arbeidstid.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/bullshitbench_health.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/bullshitbench_v1_v2.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/health.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/hei_refusal.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/helfo.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/helpmed.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/human_rights_education.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/human_rights_fair_trial.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/human_rights_water.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/blank_chart.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/illegible_label.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/injection.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/make_images.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/product_a.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/product_b.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/quarterly_revenue.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/truncated_axis.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/lanekassen.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/nav_aap.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/nb_kryss_ordning.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/rag.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/safety.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/skatteetaten.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/skatteetaten_legitimasjon.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/system_prompt.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/toll_reisegodskvote.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/ung.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/vision_integrity.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/spans.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/stats.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/__init__.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/base.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/callable.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/http.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/context.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/shared.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/store.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/utils.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/scenario_viewer.html +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/server.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/thumbnail.png +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/visualizer.html +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/dependency_links.txt +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/entry_points.txt +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/requires.txt +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/top_level.txt +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_anthropic_judge_default_path.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_audit_flow.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_basic.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist_judge.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist_pipeline.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist_reframing.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_cross_judge.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_custom_prompts.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_document_expansion.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_experiment_streaming.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_fakes.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_file_uri.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_fragility.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_generation_params.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_json_format.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_compose.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_notes.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_output_contract.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_postprocess.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_registry.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_response_schema.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_robustness_pipeline.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_model_auditor.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_perturbations.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_reframing.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_rejudge.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_repeated_experiments.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_resumable_experiments.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_scenario_data.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_scenario_pack_conventions.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_spans.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_stats.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_strip_thinking.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_target_api_key.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_token_counting.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_tracing_auth.py +0 -0
- {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_vision_integrity_pack.py +0 -0
|
@@ -34,7 +34,7 @@ except PackageNotFoundError:
|
|
|
34
34
|
__version__ = "0.0.0+unknown"
|
|
35
35
|
__author__ = "SimpleAudit Contributors"
|
|
36
36
|
|
|
37
|
-
from .model_auditor import ModelAuditor
|
|
37
|
+
from .model_auditor import ModelAuditor, ScenarioExecution
|
|
38
38
|
from .auditor import Auditor
|
|
39
39
|
from .targets import (
|
|
40
40
|
CallableTarget,
|
|
@@ -86,6 +86,7 @@ from .perturbations import (
|
|
|
86
86
|
|
|
87
87
|
__all__ = [
|
|
88
88
|
"ModelAuditor",
|
|
89
|
+
"ScenarioExecution",
|
|
89
90
|
"Auditor",
|
|
90
91
|
"Target",
|
|
91
92
|
"TargetContext",
|
|
@@ -43,6 +43,31 @@ from .utils import (
|
|
|
43
43
|
from .utils import parse_json_response as _parse_json_response
|
|
44
44
|
|
|
45
45
|
|
|
46
|
+
from dataclasses import dataclass
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
@dataclass
|
|
50
|
+
class ScenarioExecution:
|
|
51
|
+
"""Captures the state of a scenario after target execution but before judgment.
|
|
52
|
+
|
|
53
|
+
This dataclass represents one scenario execution through the target conversation phase.
|
|
54
|
+
It holds all state necessary for the evidence resolver to select trace evidence, which
|
|
55
|
+
can then be passed to the judge.
|
|
56
|
+
"""
|
|
57
|
+
scenario_name: str
|
|
58
|
+
scenario_description: str
|
|
59
|
+
conversation: List[Dict]
|
|
60
|
+
trace_correlation: Optional[Any] = None
|
|
61
|
+
audit_run_id: Optional[str] = None
|
|
62
|
+
error: Optional[str] = None
|
|
63
|
+
scenario_meta: Optional[Dict[str, Any]] = None
|
|
64
|
+
expected_behavior: Optional[List[str]] = None
|
|
65
|
+
target_input_tokens: int = 0
|
|
66
|
+
target_output_tokens: int = 0
|
|
67
|
+
auditor_input_tokens: int = 0
|
|
68
|
+
auditor_output_tokens: int = 0
|
|
69
|
+
|
|
70
|
+
|
|
46
71
|
def _user_agent() -> str:
|
|
47
72
|
try:
|
|
48
73
|
return f"simpleaudit/{_pkg_version('simpleaudit')}"
|
|
@@ -867,6 +892,7 @@ Evaluate this conversation and respond with this exact JSON structure:
|
|
|
867
892
|
evidence_spans: Optional[List[Dict[str, Any]]] = None,
|
|
868
893
|
audit_run_id: Optional[str] = None,
|
|
869
894
|
trace_correlation: Optional[Any] = None,
|
|
895
|
+
evidence_resolver: Optional[Callable[[ScenarioExecution], Union[List[Dict[str, Any]], None]]] = None,
|
|
870
896
|
) -> AuditResult:
|
|
871
897
|
turns = max_turns or self.max_turns
|
|
872
898
|
# Per-scenario correlation ids. A fresh trace id per scenario keeps each
|
|
@@ -976,6 +1002,30 @@ Evaluate this conversation and respond with this exact JSON structure:
|
|
|
976
1002
|
self._log(f"--- Scenario FAILED mid-conversation: {name} [{error}] ---")
|
|
977
1003
|
|
|
978
1004
|
if error is None:
|
|
1005
|
+
# Call evidence resolver if provided, after target conversation, before judge.
|
|
1006
|
+
resolved_evidence = evidence_spans
|
|
1007
|
+
if evidence_resolver is not None:
|
|
1008
|
+
execution = ScenarioExecution(
|
|
1009
|
+
scenario_name=name,
|
|
1010
|
+
scenario_description=description,
|
|
1011
|
+
conversation=conversation,
|
|
1012
|
+
scenario_meta=scenario_meta,
|
|
1013
|
+
expected_behavior=expected_behavior,
|
|
1014
|
+
trace_correlation=trace_correlation,
|
|
1015
|
+
audit_run_id=audit_run_id,
|
|
1016
|
+
error=None,
|
|
1017
|
+
target_input_tokens=target_input_tokens,
|
|
1018
|
+
target_output_tokens=target_output_tokens,
|
|
1019
|
+
auditor_input_tokens=auditor_input_tokens,
|
|
1020
|
+
auditor_output_tokens=auditor_output_tokens,
|
|
1021
|
+
)
|
|
1022
|
+
try:
|
|
1023
|
+
resolved_evidence = await evidence_resolver(execution) if asyncio.iscoroutinefunction(evidence_resolver) else evidence_resolver(execution)
|
|
1024
|
+
except Exception as exc:
|
|
1025
|
+
self._log(f"--- Evidence resolution failed: {type(exc).__name__}: {exc} ---")
|
|
1026
|
+
# If evidence resolution fails, continue without evidence
|
|
1027
|
+
resolved_evidence = evidence_spans
|
|
1028
|
+
|
|
979
1029
|
self._log("Judging conversation...", name=name)
|
|
980
1030
|
judge_prompt, judge_schema, judge_postprocess, fell_back = self._resolve_judge_spec(
|
|
981
1031
|
self.judge_prompt,
|
|
@@ -1003,7 +1053,7 @@ Evaluate this conversation and respond with this exact JSON structure:
|
|
|
1003
1053
|
postprocess=judge_postprocess,
|
|
1004
1054
|
scenario_meta=scenario_meta,
|
|
1005
1055
|
params=effective_judge or None,
|
|
1006
|
-
evidence_spans=
|
|
1056
|
+
evidence_spans=resolved_evidence,
|
|
1007
1057
|
)
|
|
1008
1058
|
judge_input_tokens += j_in
|
|
1009
1059
|
judge_output_tokens += j_out
|
|
@@ -1076,6 +1126,7 @@ Evaluate this conversation and respond with this exact JSON structure:
|
|
|
1076
1126
|
evidence_spans: Optional[List[Dict[str, Any]]] = None,
|
|
1077
1127
|
audit_run_id: Optional[str] = None,
|
|
1078
1128
|
trace_correlation: Optional[Any] = None,
|
|
1129
|
+
evidence_resolver: Optional[Callable[[ScenarioExecution], Union[List[Dict[str, Any]], None]]] = None,
|
|
1079
1130
|
) -> AuditResults:
|
|
1080
1131
|
if max_workers < 1:
|
|
1081
1132
|
raise ValueError(
|
|
@@ -1150,6 +1201,7 @@ Evaluate this conversation and respond with this exact JSON structure:
|
|
|
1150
1201
|
evidence_spans=evidence_spans,
|
|
1151
1202
|
audit_run_id=audit_run_id,
|
|
1152
1203
|
trace_correlation=trace_correlation,
|
|
1204
|
+
evidence_resolver=evidence_resolver,
|
|
1153
1205
|
)
|
|
1154
1206
|
except Exception as exc:
|
|
1155
1207
|
# Don't let one failing scenario abort the whole batch and
|
|
@@ -1210,6 +1262,7 @@ Evaluate this conversation and respond with this exact JSON structure:
|
|
|
1210
1262
|
evidence_spans: Optional[List[Dict[str, Any]]] = None,
|
|
1211
1263
|
audit_run_id: Optional[str] = None,
|
|
1212
1264
|
trace_correlation: Optional[Any] = None,
|
|
1265
|
+
evidence_resolver: Optional[Callable[[ScenarioExecution], Union[List[Dict[str, Any]], None]]] = None,
|
|
1213
1266
|
) -> AuditResults:
|
|
1214
1267
|
try:
|
|
1215
1268
|
asyncio.get_running_loop()
|
|
@@ -1228,6 +1281,7 @@ Evaluate this conversation and respond with this exact JSON structure:
|
|
|
1228
1281
|
evidence_spans=evidence_spans,
|
|
1229
1282
|
audit_run_id=audit_run_id,
|
|
1230
1283
|
trace_correlation=trace_correlation,
|
|
1284
|
+
evidence_resolver=evidence_resolver,
|
|
1231
1285
|
)
|
|
1232
1286
|
)
|
|
1233
1287
|
msg = "ModelAuditor.run() cannot be called from an active event loop. Use await <object>.run_async()."
|
|
@@ -88,6 +88,7 @@ class ModelTarget:
|
|
|
88
88
|
documents=documents,
|
|
89
89
|
response_format=response_format,
|
|
90
90
|
params=params,
|
|
91
|
+
context=context,
|
|
91
92
|
max_retries=self.max_retries,
|
|
92
93
|
retry_backoff=self.retry_backoff,
|
|
93
94
|
)
|
|
@@ -104,6 +105,7 @@ async def _client_send(
|
|
|
104
105
|
documents: Optional[List[Union[str, Dict[str, Any]]]],
|
|
105
106
|
response_format: Optional[Dict[str, Any]],
|
|
106
107
|
params: Optional[Dict[str, Any]],
|
|
108
|
+
context: Optional[TargetContext] = None,
|
|
107
109
|
max_retries: int,
|
|
108
110
|
retry_backoff: float,
|
|
109
111
|
) -> TargetResponse:
|
|
@@ -111,9 +113,21 @@ async def _client_send(
|
|
|
111
113
|
|
|
112
114
|
Imported lazily so that ``simpleaudit.targets`` does not create an import
|
|
113
115
|
cycle with ``model_auditor`` at package import time.
|
|
116
|
+
|
|
117
|
+
When the per-turn context carries W3C trace headers (e.g. the engine's
|
|
118
|
+
scenario traceparent), they are merged into ``extra_headers`` so the
|
|
119
|
+
target's outgoing call propagates the trace context. This is what lets an
|
|
120
|
+
instrumented target export spans under the same trace id the engine's
|
|
121
|
+
TraceCorrelation records. User-supplied extra_headers are preserved.
|
|
114
122
|
"""
|
|
115
123
|
from ..model_auditor import ModelAuditor
|
|
116
124
|
|
|
125
|
+
if context is not None and context.trace_headers:
|
|
126
|
+
headers = dict(params.get("extra_headers", {}) if params else {})
|
|
127
|
+
headers.update(context.trace_headers)
|
|
128
|
+
params = dict(params or {})
|
|
129
|
+
params["extra_headers"] = headers
|
|
130
|
+
|
|
117
131
|
content, input_tokens, output_tokens = await ModelAuditor._call_async(
|
|
118
132
|
client,
|
|
119
133
|
model,
|
|
@@ -15,11 +15,14 @@ when the target is instrumented.
|
|
|
15
15
|
from .auth import (
|
|
16
16
|
AuthResult,
|
|
17
17
|
Authenticator,
|
|
18
|
+
OTLPAuth,
|
|
19
|
+
TARGET_ID_ATTR,
|
|
18
20
|
hash_secret,
|
|
19
21
|
make_basic_bearer_authenticator,
|
|
20
22
|
new_salt,
|
|
21
23
|
parse_basic_header,
|
|
22
24
|
parse_bearer_header,
|
|
25
|
+
tag_spans,
|
|
23
26
|
token_lookup_prefix,
|
|
24
27
|
verify_secret,
|
|
25
28
|
)
|
|
@@ -71,6 +74,9 @@ __all__ = [
|
|
|
71
74
|
"DEFAULT_NOISE_KINDS",
|
|
72
75
|
"AuthResult",
|
|
73
76
|
"Authenticator",
|
|
77
|
+
"OTLPAuth",
|
|
78
|
+
"TARGET_ID_ATTR",
|
|
79
|
+
"tag_spans",
|
|
74
80
|
"hash_secret",
|
|
75
81
|
"new_salt",
|
|
76
82
|
"verify_secret",
|
|
@@ -0,0 +1,332 @@
|
|
|
1
|
+
"""
|
|
2
|
+
OTLP receiver authentication — the "WHO is allowed to submit" layer.
|
|
3
|
+
|
|
4
|
+
This module is deliberately **framework-agnostic**: it has no Django (or any
|
|
5
|
+
web-framework) dependency so both the standalone ``simpleaudit`` library and
|
|
6
|
+
the SimpleAuditStudio app can use the same auth. It answers one question —
|
|
7
|
+
which target is allowed to push spans — and maps the credential to a stable
|
|
8
|
+
``target_id`` that the receiver tags onto every ingested span.
|
|
9
|
+
|
|
10
|
+
Two auth modes (plus ``none``):
|
|
11
|
+
|
|
12
|
+
- ``none`` — no check. Local / ephemeral receiver on a trusted network.
|
|
13
|
+
- ``basic`` — ``Authorization: Basic base64(username:password)``.
|
|
14
|
+
Best for OpenWebUI, which natively supports
|
|
15
|
+
``OTEL_BASIC_AUTH_USERNAME`` / ``OTEL_BASIC_AUTH_PASSWORD``.
|
|
16
|
+
- ``bearer`` — ``Authorization: Bearer <token>``. Best for generic OTel
|
|
17
|
+
exporters that can set arbitrary headers
|
|
18
|
+
(``OTEL_EXPORTER_OTLP_HEADERS=Authorization=Bearer <t>``).
|
|
19
|
+
|
|
20
|
+
Credentials are revocable and rotatable. Secrets are stored as salted
|
|
21
|
+
SHA-256 hashes and compared in constant time (``hmac.compare_digest``) to
|
|
22
|
+
avoid timing side-channels.
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
from __future__ import annotations
|
|
26
|
+
|
|
27
|
+
import base64
|
|
28
|
+
import binascii
|
|
29
|
+
import hashlib
|
|
30
|
+
import hmac
|
|
31
|
+
import os
|
|
32
|
+
import secrets
|
|
33
|
+
from dataclasses import dataclass
|
|
34
|
+
from typing import Callable, Dict, Optional, Protocol, Tuple
|
|
35
|
+
|
|
36
|
+
# Salt length (bytes) prepended to each secret before hashing.
|
|
37
|
+
_SALT_LEN = 16
|
|
38
|
+
# Attribute key the receiver stamps onto every authenticated span.
|
|
39
|
+
TARGET_ID_ATTR = "simpleaudit.target_id"
|
|
40
|
+
|
|
41
|
+
SALT_LEN = _SALT_LEN
|
|
42
|
+
BEARER_TOKEN_PREFIX = "sa_otlp_"
|
|
43
|
+
TOKEN_LOOKUP_PREFIX_LEN = 16
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def new_salt() -> bytes:
|
|
47
|
+
return _new_salt()
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def hash_secret(secret: str, salt: bytes) -> bytes:
|
|
51
|
+
return _hash_secret(secret, salt)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def verify_secret(secret: str, salt: bytes, expected_hash: bytes) -> bool:
|
|
55
|
+
if not salt or not expected_hash:
|
|
56
|
+
return False
|
|
57
|
+
return hmac.compare_digest(expected_hash, _hash_secret(secret, salt))
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def generate_password() -> str:
|
|
61
|
+
return secrets.token_urlsafe(32)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def generate_token() -> str:
|
|
65
|
+
return BEARER_TOKEN_PREFIX + secrets.token_urlsafe(32)
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def token_lookup_prefix(token: str) -> str:
|
|
69
|
+
return (token or "")[:TOKEN_LOOKUP_PREFIX_LEN]
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def parse_basic_header(authorization: Optional[str]) -> Optional[Tuple[str, str]]:
|
|
73
|
+
if not authorization:
|
|
74
|
+
return None
|
|
75
|
+
parts = authorization.strip().split(" ", 1)
|
|
76
|
+
if len(parts) != 2 or parts[0].lower() != "basic":
|
|
77
|
+
return None
|
|
78
|
+
try:
|
|
79
|
+
decoded = base64.b64decode(parts[1].strip(), validate=True).decode("utf-8")
|
|
80
|
+
except (binascii.Error, UnicodeDecodeError, ValueError):
|
|
81
|
+
return None
|
|
82
|
+
if ":" not in decoded:
|
|
83
|
+
return None
|
|
84
|
+
return tuple(decoded.split(":", 1))
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def parse_bearer_header(authorization: Optional[str]) -> Optional[str]:
|
|
88
|
+
if not authorization:
|
|
89
|
+
return None
|
|
90
|
+
parts = authorization.strip().split(" ", 1)
|
|
91
|
+
if len(parts) != 2 or parts[0].lower() != "bearer":
|
|
92
|
+
return None
|
|
93
|
+
return parts[1].strip() or None
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class Authenticator(Protocol):
|
|
97
|
+
def __call__(self, authorization: Optional[str]) -> "AuthResult": ...
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def make_basic_bearer_authenticator(
|
|
101
|
+
verify: Callable[[Optional[str]], Optional[str]],
|
|
102
|
+
) -> Authenticator:
|
|
103
|
+
def _authenticate(authorization: Optional[str]) -> AuthResult:
|
|
104
|
+
identity = verify(authorization)
|
|
105
|
+
return AuthResult(ok=identity is not None, target_id=identity)
|
|
106
|
+
|
|
107
|
+
return _authenticate
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def _hash_secret(secret: str, salt: bytes) -> bytes:
|
|
111
|
+
"""Salted SHA-256 of a secret. Deterministic for a given salt."""
|
|
112
|
+
return hashlib.sha256(salt + secret.encode("utf-8")).digest()
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def _new_salt() -> bytes:
|
|
116
|
+
return os.urandom(_SALT_LEN)
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
@dataclass
|
|
120
|
+
class _Credential:
|
|
121
|
+
"""A stored credential: salt + hash of the secret, bound to a target_id."""
|
|
122
|
+
|
|
123
|
+
target_id: str
|
|
124
|
+
salt: bytes
|
|
125
|
+
secret_hash: bytes
|
|
126
|
+
revoked: bool = False
|
|
127
|
+
|
|
128
|
+
def matches(self, secret: str) -> bool:
|
|
129
|
+
"""Constant-time check of a presented secret against the stored hash."""
|
|
130
|
+
if self.revoked:
|
|
131
|
+
return False
|
|
132
|
+
return hmac.compare_digest(self.secret_hash, _hash_secret(secret, self.salt))
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
@dataclass(init=False)
|
|
136
|
+
class AuthResult:
|
|
137
|
+
"""Outcome of an :meth:`OTLPAuth.verify` call.
|
|
138
|
+
|
|
139
|
+
``ok`` is True when the request may proceed. ``target_id`` is set (and
|
|
140
|
+
non-empty) when a credential was matched; it is ``None`` in ``none`` mode
|
|
141
|
+
and on failure. ``reason`` is a short, log-safe explanation on failure.
|
|
142
|
+
"""
|
|
143
|
+
|
|
144
|
+
ok: bool
|
|
145
|
+
target_id: Optional[str] = None
|
|
146
|
+
reason: Optional[str] = None
|
|
147
|
+
|
|
148
|
+
def __init__(
|
|
149
|
+
self,
|
|
150
|
+
ok: Optional[bool] = None,
|
|
151
|
+
target_id: Optional[str] = None,
|
|
152
|
+
reason: Optional[str] = None,
|
|
153
|
+
*,
|
|
154
|
+
authenticated: Optional[bool] = None,
|
|
155
|
+
identity: Optional[str] = None,
|
|
156
|
+
) -> None:
|
|
157
|
+
self.ok = bool(ok if ok is not None else authenticated)
|
|
158
|
+
self.target_id = target_id if target_id is not None else identity
|
|
159
|
+
self.reason = reason
|
|
160
|
+
|
|
161
|
+
@property
|
|
162
|
+
def authenticated(self) -> bool:
|
|
163
|
+
return self.ok
|
|
164
|
+
|
|
165
|
+
@property
|
|
166
|
+
def identity(self) -> Optional[str]:
|
|
167
|
+
return self.target_id
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
class OTLPAuth:
|
|
171
|
+
"""Credential store + verifier for an OTLP receiver.
|
|
172
|
+
|
|
173
|
+
Construct with a mode, add credentials, then hand the instance to an
|
|
174
|
+
:class:`~simpleaudit.tracing.otlp.OTLPTraceReceiver` or
|
|
175
|
+
:class:`~simpleaudit.tracing.otlp.EphemeralOTLPReceiver`. The receiver
|
|
176
|
+
calls :meth:`verify` on each request's ``Authorization`` header.
|
|
177
|
+
|
|
178
|
+
Example::
|
|
179
|
+
|
|
180
|
+
auth = OTLPAuth(mode="basic")
|
|
181
|
+
auth.add_basic("sa_owui_17", "p4ss", target_id="owui_17")
|
|
182
|
+
result = auth.verify("Basic " + b64("sa_owui_17:p4ss"))
|
|
183
|
+
assert result.ok and result.target_id == "owui_17"
|
|
184
|
+
"""
|
|
185
|
+
|
|
186
|
+
def __init__(self, mode: str = "none") -> None:
|
|
187
|
+
mode = (mode or "none").strip().lower()
|
|
188
|
+
if mode not in ("none", "basic", "bearer"):
|
|
189
|
+
raise ValueError(f"Unknown OTLPAuth mode: {mode!r} (expected 'none', 'basic', or 'bearer')")
|
|
190
|
+
self.mode = mode
|
|
191
|
+
# username -> credential (basic)
|
|
192
|
+
self._basic: Dict[str, _Credential] = {}
|
|
193
|
+
# token-hash-keyed credential (bearer). We index by a *lookup* hash of
|
|
194
|
+
# the token (salted per-credential) so we can find the record to
|
|
195
|
+
# compare against without storing the raw token.
|
|
196
|
+
self._bearer: Dict[str, _Credential] = {}
|
|
197
|
+
|
|
198
|
+
# ── Credential management ────────────────────────────────────────────
|
|
199
|
+
|
|
200
|
+
def add_basic(self, username: str, password: str, target_id: str) -> None:
|
|
201
|
+
"""Register (or replace) a Basic-Auth credential for ``target_id``."""
|
|
202
|
+
if not username:
|
|
203
|
+
raise ValueError("username is required")
|
|
204
|
+
salt = _new_salt()
|
|
205
|
+
self._basic[username] = _Credential(
|
|
206
|
+
target_id=target_id, salt=salt, secret_hash=_hash_secret(password, salt)
|
|
207
|
+
)
|
|
208
|
+
|
|
209
|
+
def add_bearer(self, token: str, target_id: str) -> None:
|
|
210
|
+
"""Register (or replace) a Bearer credential for ``target_id``."""
|
|
211
|
+
if not token:
|
|
212
|
+
raise ValueError("token is required")
|
|
213
|
+
salt = _new_salt()
|
|
214
|
+
cred = _Credential(target_id=target_id, salt=salt, secret_hash=_hash_secret(token, salt))
|
|
215
|
+
# Index by the salted hash so lookup needs only the presented token.
|
|
216
|
+
self._bearer[_hash_secret(token, salt).hex()] = cred
|
|
217
|
+
|
|
218
|
+
def revoke(self, username_or_token: str) -> bool:
|
|
219
|
+
"""Revoke a credential by username (basic) or token (bearer).
|
|
220
|
+
|
|
221
|
+
Returns True if a matching credential was found and revoked.
|
|
222
|
+
"""
|
|
223
|
+
if username_or_token in self._basic:
|
|
224
|
+
self._basic[username_or_token].revoked = True
|
|
225
|
+
return True
|
|
226
|
+
# For bearer we must find the record by hashing the token with each
|
|
227
|
+
# stored salt; the index key is that hash, so we can't reverse it.
|
|
228
|
+
# Instead we scan (credential sets are small — one per target).
|
|
229
|
+
for cred in self._bearer.values():
|
|
230
|
+
if cred.matches(username_or_token):
|
|
231
|
+
cred.revoked = True
|
|
232
|
+
return True
|
|
233
|
+
return False
|
|
234
|
+
|
|
235
|
+
def rotate_basic(self, username: str, new_password: str) -> None:
|
|
236
|
+
"""Re-issue the password for an existing Basic credential (same target)."""
|
|
237
|
+
cred = self._basic.get(username)
|
|
238
|
+
if cred is None:
|
|
239
|
+
raise KeyError(f"No Basic credential for username {username!r}")
|
|
240
|
+
salt = _new_salt()
|
|
241
|
+
cred.salt = salt
|
|
242
|
+
cred.secret_hash = _hash_secret(new_password, salt)
|
|
243
|
+
cred.revoked = False
|
|
244
|
+
|
|
245
|
+
def rotate_bearer(self, old_token: str, new_token: str) -> None:
|
|
246
|
+
"""Replace a Bearer token, keeping the same target_id."""
|
|
247
|
+
old_cred = self._find_bearer(old_token)
|
|
248
|
+
if old_cred is None:
|
|
249
|
+
raise KeyError("No Bearer credential for the given token")
|
|
250
|
+
# Remove the old index entry and register the new token under the
|
|
251
|
+
# same target_id.
|
|
252
|
+
for key, cred in list(self._bearer.items()):
|
|
253
|
+
if cred is old_cred:
|
|
254
|
+
del self._bearer[key]
|
|
255
|
+
break
|
|
256
|
+
self.add_bearer(new_token, old_cred.target_id)
|
|
257
|
+
|
|
258
|
+
def _find_bearer(self, token: str) -> Optional[_Credential]:
|
|
259
|
+
for cred in self._bearer.values():
|
|
260
|
+
if cred.matches(token):
|
|
261
|
+
return cred
|
|
262
|
+
return None
|
|
263
|
+
|
|
264
|
+
# ── Verification ─────────────────────────────────────────────────────
|
|
265
|
+
|
|
266
|
+
def verify(self, authorization_header: Optional[str]) -> AuthResult:
|
|
267
|
+
"""Verify an ``Authorization`` header and return the target id.
|
|
268
|
+
|
|
269
|
+
In ``none`` mode this always succeeds (``target_id`` is ``None``).
|
|
270
|
+
"""
|
|
271
|
+
if self.mode == "none":
|
|
272
|
+
return AuthResult(ok=True)
|
|
273
|
+
|
|
274
|
+
header = (authorization_header or "").strip()
|
|
275
|
+
if not header:
|
|
276
|
+
return AuthResult(ok=False, reason="missing_authorization_header")
|
|
277
|
+
|
|
278
|
+
parts = header.split(" ", 1)
|
|
279
|
+
if len(parts) != 2 or not parts[1].strip():
|
|
280
|
+
return AuthResult(ok=False, reason="malformed_authorization_header")
|
|
281
|
+
scheme, value = parts[0].strip().lower(), parts[1].strip()
|
|
282
|
+
|
|
283
|
+
# The scheme must match the configured mode.
|
|
284
|
+
if self.mode == "basic":
|
|
285
|
+
if scheme != "basic":
|
|
286
|
+
return AuthResult(ok=False, reason="expected_basic_scheme")
|
|
287
|
+
return self._verify_basic(value)
|
|
288
|
+
# bearer
|
|
289
|
+
if scheme != "bearer":
|
|
290
|
+
return AuthResult(ok=False, reason="expected_bearer_scheme")
|
|
291
|
+
return self._verify_bearer(value)
|
|
292
|
+
|
|
293
|
+
def _verify_basic(self, value: str) -> AuthResult:
|
|
294
|
+
# value should be base64("username:password")
|
|
295
|
+
try:
|
|
296
|
+
decoded = base64.b64decode(value, validate=True).decode("utf-8")
|
|
297
|
+
except (binascii.Error, UnicodeDecodeError, ValueError):
|
|
298
|
+
return AuthResult(ok=False, reason="invalid_basic_encoding")
|
|
299
|
+
if ":" not in decoded:
|
|
300
|
+
return AuthResult(ok=False, reason="invalid_basic_credentials")
|
|
301
|
+
username, password = decoded.split(":", 1)
|
|
302
|
+
cred = self._basic.get(username)
|
|
303
|
+
if cred is None:
|
|
304
|
+
return AuthResult(ok=False, reason="unknown_username")
|
|
305
|
+
if not cred.matches(password):
|
|
306
|
+
return AuthResult(ok=False, reason="invalid_credentials")
|
|
307
|
+
return AuthResult(ok=True, target_id=cred.target_id)
|
|
308
|
+
|
|
309
|
+
def _verify_bearer(self, token: str) -> AuthResult:
|
|
310
|
+
if not token:
|
|
311
|
+
return AuthResult(ok=False, reason="empty_bearer_token")
|
|
312
|
+
cred = self._find_bearer(token)
|
|
313
|
+
if cred is None:
|
|
314
|
+
return AuthResult(ok=False, reason="unknown_token")
|
|
315
|
+
return AuthResult(ok=True, target_id=cred.target_id)
|
|
316
|
+
|
|
317
|
+
|
|
318
|
+
def tag_spans(spans, target_id: Optional[str]) -> None:
|
|
319
|
+
"""Stamp ``simpleaudit.target_id`` onto each span's attributes in place.
|
|
320
|
+
|
|
321
|
+
No-op when ``target_id`` is falsy (e.g. ``none`` mode). Mutates the raw
|
|
322
|
+
span dicts before they are normalized/stored, so the tag flows through
|
|
323
|
+
:func:`simpleaudit.tracing.otlp.parse_otlp_json` and is queryable via
|
|
324
|
+
``SpanStore.by_attribute``.
|
|
325
|
+
"""
|
|
326
|
+
if not target_id:
|
|
327
|
+
return
|
|
328
|
+
for span in spans:
|
|
329
|
+
attrs = span.get("attributes")
|
|
330
|
+
if attrs is None:
|
|
331
|
+
span["attributes"] = attrs = {}
|
|
332
|
+
attrs.setdefault(TARGET_ID_ATTR, target_id)
|
|
@@ -21,6 +21,7 @@ from __future__ import annotations
|
|
|
21
21
|
import json
|
|
22
22
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
|
23
23
|
|
|
24
|
+
from .auth import OTLPAuth, tag_spans
|
|
24
25
|
from .store import SpanStore, normalize_span
|
|
25
26
|
|
|
26
27
|
if TYPE_CHECKING: # pragma: no cover - annotation only
|
|
@@ -132,13 +133,16 @@ class OTLPTraceReceiver:
|
|
|
132
133
|
self,
|
|
133
134
|
store: Optional[SpanStore] = None,
|
|
134
135
|
authenticator: Optional["Authenticator"] = None,
|
|
136
|
+
auth: Optional[OTLPAuth] = None,
|
|
135
137
|
) -> None:
|
|
138
|
+
self._default_store = store is None
|
|
136
139
|
self.store = store or SpanStore()
|
|
137
140
|
# Optional auth gate. When set, handle() calls it with the request's
|
|
138
141
|
# Authorization header and rejects the export (401) on failure. When
|
|
139
142
|
# None the receiver is open — the historical, backward-compatible
|
|
140
143
|
# behaviour for the local, single-audit ephemeral case.
|
|
141
144
|
self.authenticator = authenticator
|
|
145
|
+
self.auth = auth
|
|
142
146
|
|
|
143
147
|
async def handle(self, body: Any, authorization: Optional[str] = None) -> Dict[str, Any]:
|
|
144
148
|
"""Ingest an OTLP JSON export body; return the OTLP ack payload.
|
|
@@ -147,14 +151,30 @@ class OTLPTraceReceiver:
|
|
|
147
151
|
``Authorization`` header) is checked first; a failed check returns a
|
|
148
152
|
401 ``{"error": ...}`` payload and no spans are stored.
|
|
149
153
|
"""
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
+
result = self._authenticate(authorization)
|
|
155
|
+
if result is not None and not result.authenticated:
|
|
156
|
+
return {
|
|
157
|
+
"status": 401,
|
|
158
|
+
"authenticated": False,
|
|
159
|
+
"reason": getattr(result, "reason", None),
|
|
160
|
+
"error": {"code": "unauthorized", "message": "Invalid or missing OTLP credentials."},
|
|
161
|
+
}
|
|
154
162
|
raw_spans = parse_otlp_json(body)
|
|
163
|
+
if result is not None:
|
|
164
|
+
tag_spans(raw_spans, result.identity)
|
|
155
165
|
self.store.add_many(raw_spans)
|
|
156
166
|
# OTLP ack: partialSuccess with the number of rejected spans (0 here).
|
|
157
|
-
|
|
167
|
+
ack = {"partialSuccess": {"rejectedSpans": 0}}
|
|
168
|
+
if self.auth is not None or self._default_store:
|
|
169
|
+
ack["authenticated"] = True
|
|
170
|
+
return ack
|
|
171
|
+
|
|
172
|
+
def _authenticate(self, authorization: Optional[str]) -> Any:
|
|
173
|
+
if self.auth is not None:
|
|
174
|
+
return self.auth.verify(authorization)
|
|
175
|
+
if self.authenticator is not None:
|
|
176
|
+
return self.authenticator(authorization)
|
|
177
|
+
return None
|
|
158
178
|
|
|
159
179
|
def trace_ids(self) -> List[str]:
|
|
160
180
|
seen: List[str] = []
|
|
@@ -196,6 +216,7 @@ class EphemeralOTLPReceiver:
|
|
|
196
216
|
port: int = 0,
|
|
197
217
|
store: Optional[SpanStore] = None,
|
|
198
218
|
authenticator: Optional["Authenticator"] = None,
|
|
219
|
+
auth: Optional[OTLPAuth] = None,
|
|
199
220
|
) -> None:
|
|
200
221
|
self.host = host
|
|
201
222
|
self.port = port
|
|
@@ -204,6 +225,7 @@ class EphemeralOTLPReceiver:
|
|
|
204
225
|
# checked against the request's Authorization header and rejected with
|
|
205
226
|
# 401 on failure. None = open receiver (default).
|
|
206
227
|
self.authenticator = authenticator
|
|
228
|
+
self.auth = auth
|
|
207
229
|
self._runner: Optional[Any] = None
|
|
208
230
|
self._site: Optional[Any] = None
|
|
209
231
|
self._thread: Optional[Any] = None
|
|
@@ -225,13 +247,12 @@ class EphemeralOTLPReceiver:
|
|
|
225
247
|
async def _handle_traces(self, request: Any) -> Any:
|
|
226
248
|
from aiohttp import web
|
|
227
249
|
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
234
|
-
)
|
|
250
|
+
result = self._authenticate(request.headers.get("Authorization"))
|
|
251
|
+
if result is not None and not result.authenticated:
|
|
252
|
+
return web.json_response(
|
|
253
|
+
{"authenticated": False, "error": {"code": "unauthorized", "message": "Invalid or missing OTLP credentials."}},
|
|
254
|
+
status=401,
|
|
255
|
+
)
|
|
235
256
|
|
|
236
257
|
content_type = request.headers.get("Content-Type", "")
|
|
237
258
|
body_bytes = await request.read()
|
|
@@ -240,12 +261,21 @@ class EphemeralOTLPReceiver:
|
|
|
240
261
|
raw_spans = _parse_otlp_http_protobuf(body_bytes)
|
|
241
262
|
else:
|
|
242
263
|
raw_spans = parse_otlp_json(body_bytes.decode("utf-8"))
|
|
264
|
+
if result is not None:
|
|
265
|
+
tag_spans(raw_spans, result.identity)
|
|
243
266
|
self.store.add_many(raw_spans)
|
|
244
267
|
except Exception:
|
|
245
268
|
# Never fail the export; ack with a rejection count so the target
|
|
246
269
|
# doesn't retry-loop. The audit continues regardless.
|
|
247
270
|
return web.json_response({"partialSuccess": {"rejectedSpans": 1}}, status=200)
|
|
248
|
-
return web.json_response({"partialSuccess": {"rejectedSpans": 0}}, status=200)
|
|
271
|
+
return web.json_response({"partialSuccess": {"rejectedSpans": 0}, "authenticated": True}, status=200)
|
|
272
|
+
|
|
273
|
+
def _authenticate(self, authorization: Optional[str]) -> Any:
|
|
274
|
+
if self.auth is not None:
|
|
275
|
+
return self.auth.verify(authorization)
|
|
276
|
+
if self.authenticator is not None:
|
|
277
|
+
return self.authenticator(authorization)
|
|
278
|
+
return None
|
|
249
279
|
|
|
250
280
|
def _serve(self, loop: Any) -> None:
|
|
251
281
|
import asyncio
|