466 lines
17 KiB
Python
466 lines
17 KiB
Python
"""Connectivity validation helpers for built-in datasource overrides."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from datetime import UTC, datetime
|
||
import hashlib
|
||
import json
|
||
import os
|
||
from typing import Any
|
||
|
||
import httpx
|
||
from sqlalchemy import func, select
|
||
|
||
from app.core.data_sources import get_data_sources_config
|
||
from app.core.datasource_defaults import DEFAULT_DATASOURCES
|
||
from app.models.collected_data import CollectedData
|
||
from app.models.datasource import DataSource
|
||
from app.models.datasource_config import DataSourceConfig
|
||
from app.models.system_setting import SystemSetting
|
||
from app.services.barentswatch import (
|
||
BarentsWatchConfig,
|
||
_read_zshrc_env,
|
||
fetch_barentswatch_access_token,
|
||
resolve_barentswatch_config,
|
||
)
|
||
|
||
|
||
CONNECTIVITY_VALIDATION_KEY = "connectivity_validation"
|
||
CONNECTIVITY_STORE_CATEGORY = "datasource_connectivity_validations"
|
||
SUPPORTED_CREDENTIAL_PROVIDERS = {"barentswatch", "spacetrack", "aisstream"}
|
||
|
||
|
||
def _sha256_json(payload: Any) -> str:
|
||
encoded = json.dumps(payload, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8")
|
||
return hashlib.sha256(encoded).hexdigest()
|
||
|
||
|
||
def _resolve_spacetrack_credentials() -> tuple[str, str, str]:
|
||
zshrc_env = _read_zshrc_env()
|
||
username = os.getenv("SPACETRACK_USERNAME") or zshrc_env.get("SPACETRACK_USERNAME") or ""
|
||
password = os.getenv("SPACETRACK_PASSWORD") or zshrc_env.get("SPACETRACK_PASSWORD") or ""
|
||
source = "environment" if os.getenv("SPACETRACK_USERNAME") or os.getenv("SPACETRACK_PASSWORD") else ""
|
||
if not source and (username or password):
|
||
source = "~/.zshrc"
|
||
return username, password, source or "missing"
|
||
|
||
|
||
def _resolve_spacetrack_credentials_with_override(
|
||
credential_override: dict[str, str] | None = None,
|
||
) -> tuple[str, str, str]:
|
||
if credential_override and (
|
||
credential_override.get("username") or credential_override.get("password")
|
||
):
|
||
return (
|
||
str(credential_override.get("username") or ""),
|
||
str(credential_override.get("password") or ""),
|
||
"draft",
|
||
)
|
||
return _resolve_spacetrack_credentials()
|
||
|
||
|
||
async def _resolve_aisstream_api_key(
|
||
db=None,
|
||
credential_override: dict[str, str] | None = None,
|
||
) -> tuple[str, str]:
|
||
if credential_override and credential_override.get("api_key"):
|
||
return str(credential_override["api_key"]), "draft"
|
||
|
||
env_key = os.getenv("AISSTREAM_API_KEY")
|
||
zshrc_key = _read_zshrc_env().get("AISSTREAM_API_KEY")
|
||
if db is not None:
|
||
result = await db.execute(
|
||
select(DataSourceConfig)
|
||
.where(DataSourceConfig.name == "aisstream_vessels")
|
||
.where(DataSourceConfig.is_active.is_(True))
|
||
)
|
||
record = result.scalar_one_or_none()
|
||
if record:
|
||
auth_config = record.auth_config or {}
|
||
runtime_config = record.config or {}
|
||
api_key = auth_config.get("api_key") or runtime_config.get("api_key")
|
||
if api_key:
|
||
return str(api_key), "datasource_config"
|
||
|
||
if env_key:
|
||
return env_key, "environment"
|
||
if zshrc_key:
|
||
return zshrc_key, "~/.zshrc"
|
||
return "", "missing"
|
||
|
||
|
||
def strip_connectivity_validation(config: dict | None) -> dict:
|
||
cleaned = dict(config or {})
|
||
cleaned.pop(CONNECTIVITY_VALIDATION_KEY, None)
|
||
return cleaned
|
||
|
||
|
||
def merge_connectivity_validation(existing_config: dict | None, next_config: dict | None) -> dict:
|
||
merged = strip_connectivity_validation(next_config)
|
||
validation = (existing_config or {}).get(CONNECTIVITY_VALIDATION_KEY)
|
||
if validation:
|
||
merged[CONNECTIVITY_VALIDATION_KEY] = validation
|
||
return merged
|
||
|
||
|
||
def get_connectivity_validation(config: DataSourceConfig | None) -> dict | None:
|
||
validation = (config.config or {}).get(CONNECTIVITY_VALIDATION_KEY) if config else None
|
||
return validation if isinstance(validation, dict) else None
|
||
|
||
|
||
async def build_builtin_connectivity_checksum(
|
||
source: str,
|
||
endpoint: str,
|
||
auth_type: str,
|
||
headers: dict | None,
|
||
config: dict | None,
|
||
db=None,
|
||
credential_override: dict[str, str] | None = None,
|
||
) -> tuple[str, dict[str, Any]]:
|
||
defaults = DEFAULT_DATASOURCES.get(source, {})
|
||
credential_provider = defaults.get("credential_provider")
|
||
credential_fingerprint = ""
|
||
credential_source = "none"
|
||
has_credentials = not defaults.get("requires_credentials", False)
|
||
|
||
if credential_provider == "barentswatch":
|
||
if credential_override:
|
||
client_id = credential_override.get("client_id", "")
|
||
client_secret = credential_override.get("client_secret", "")
|
||
credential_source = "draft"
|
||
else:
|
||
barentswatch_config = await resolve_barentswatch_config(db)
|
||
client_id = barentswatch_config.client_id
|
||
client_secret = barentswatch_config.client_secret
|
||
credential_source = barentswatch_config.credential_source
|
||
has_credentials = bool(client_id and client_secret)
|
||
credential_fingerprint = _sha256_json(
|
||
{
|
||
"client_id": client_id,
|
||
"client_secret": client_secret,
|
||
}
|
||
)
|
||
elif credential_provider == "spacetrack":
|
||
username, password, credential_source = _resolve_spacetrack_credentials_with_override(
|
||
credential_override
|
||
)
|
||
has_credentials = bool(username and password)
|
||
credential_fingerprint = _sha256_json(
|
||
{
|
||
"username": username,
|
||
"password": password,
|
||
}
|
||
)
|
||
elif credential_provider == "aisstream":
|
||
api_key, credential_source = await _resolve_aisstream_api_key(db, credential_override)
|
||
has_credentials = bool(api_key)
|
||
credential_fingerprint = _sha256_json({"api_key": api_key})
|
||
elif defaults.get("requires_credentials"):
|
||
credential_source = str(credential_provider or "unsupported")
|
||
|
||
checksum_payload = {
|
||
"source": source,
|
||
"endpoint": endpoint,
|
||
"auth_type": "none",
|
||
"headers": headers or {},
|
||
"config": strip_connectivity_validation(config),
|
||
"credential_provider": credential_provider or "none",
|
||
"credential_fingerprint": credential_fingerprint,
|
||
}
|
||
return _sha256_json(checksum_payload), {
|
||
"requires_credentials": bool(defaults.get("requires_credentials", False)),
|
||
"credential_provider": credential_provider,
|
||
"credential_source": credential_source,
|
||
"has_credentials": has_credentials,
|
||
}
|
||
|
||
|
||
async def test_builtin_connectivity(
|
||
source: str,
|
||
endpoint: str,
|
||
auth_type: str,
|
||
headers: dict | None,
|
||
config: dict | None,
|
||
db=None,
|
||
credential_override: dict[str, str] | None = None,
|
||
) -> dict[str, Any]:
|
||
defaults = DEFAULT_DATASOURCES.get(source)
|
||
if not defaults:
|
||
return {
|
||
"success": False,
|
||
"message": "未知内置采集器，无法执行连接校验。",
|
||
}
|
||
|
||
checksum, credential_context = await build_builtin_connectivity_checksum(
|
||
source,
|
||
endpoint,
|
||
auth_type,
|
||
headers,
|
||
config,
|
||
db,
|
||
credential_override,
|
||
)
|
||
if credential_context["requires_credentials"] and not credential_context["has_credentials"]:
|
||
return {
|
||
"success": False,
|
||
"checksum": checksum,
|
||
"stage": "credentials",
|
||
"message": "该采集器需要凭证，请先到采集器凭证设置中配置。",
|
||
"settings_tab": "collector_credentials",
|
||
**credential_context,
|
||
}
|
||
if (
|
||
credential_context["requires_credentials"]
|
||
and credential_context["credential_provider"] not in SUPPORTED_CREDENTIAL_PROVIDERS
|
||
):
|
||
return {
|
||
"success": False,
|
||
"checksum": checksum,
|
||
"stage": "credentials",
|
||
"message": "该采集器的凭证链路尚未接入，暂时无法完成连接校验。",
|
||
"settings_tab": "collector_credentials",
|
||
**credential_context,
|
||
}
|
||
|
||
request_headers = {str(key): str(value) for key, value in (headers or {}).items()}
|
||
request_config = strip_connectivity_validation(config)
|
||
timeout = float(request_config.get("timeout") or 30)
|
||
request_endpoint = endpoint
|
||
|
||
if credential_context["credential_provider"] == "aisstream":
|
||
if not str(request_endpoint).startswith(("ws://", "wss://")):
|
||
return {
|
||
"success": False,
|
||
"checksum": checksum,
|
||
"stage": "endpoint",
|
||
"message": "AISStream endpoint 必须是 ws:// 或 wss:// WebSocket 地址。",
|
||
**credential_context,
|
||
}
|
||
return {
|
||
"success": True,
|
||
"checksum": checksum,
|
||
"stage": "credentials",
|
||
"message": "AISStream 凭证已配置，WebSocket endpoint 格式有效。",
|
||
**credential_context,
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
|
||
if credential_context["credential_provider"] == "barentswatch":
|
||
if credential_override:
|
||
barentswatch_config = BarentsWatchConfig(
|
||
endpoint=str(request_endpoint or ""),
|
||
client_id=str(credential_override.get("client_id") or ""),
|
||
client_secret=str(credential_override.get("client_secret") or ""),
|
||
credential_source="draft",
|
||
endpoint_source="draft",
|
||
)
|
||
else:
|
||
barentswatch_config = await resolve_barentswatch_config(db)
|
||
token = await fetch_barentswatch_access_token(client, barentswatch_config)
|
||
if not token:
|
||
return {
|
||
"success": False,
|
||
"checksum": checksum,
|
||
"stage": "token",
|
||
"message": "凭证可读取，但 token 响应中没有 access_token。",
|
||
"settings_tab": "collector_credentials",
|
||
**credential_context,
|
||
}
|
||
request_headers["Authorization"] = f"Bearer {token}"
|
||
elif credential_context["credential_provider"] == "spacetrack":
|
||
username, password, _source = _resolve_spacetrack_credentials_with_override(
|
||
credential_override
|
||
)
|
||
login_url = "https://www.space-track.org/ajaxauth/login"
|
||
login_response = await client.post(
|
||
login_url,
|
||
data={
|
||
"identity": username,
|
||
"password": password,
|
||
},
|
||
)
|
||
login_response.raise_for_status()
|
||
|
||
started = datetime.now(UTC)
|
||
async with client.stream("GET", request_endpoint, headers=request_headers) as response:
|
||
response.raise_for_status()
|
||
status_code = response.status_code
|
||
elapsed_ms = (datetime.now(UTC) - started).total_seconds() * 1000
|
||
|
||
return {
|
||
"success": True,
|
||
"checksum": checksum,
|
||
"stage": "endpoint",
|
||
"message": "连接验证成功。",
|
||
"status_code": status_code,
|
||
"response_time_ms": elapsed_ms,
|
||
**credential_context,
|
||
}
|
||
except httpx.HTTPStatusError as exc:
|
||
return {
|
||
"success": False,
|
||
"checksum": checksum,
|
||
"stage": "endpoint",
|
||
"message": f"连接验证失败：HTTP {exc.response.status_code}",
|
||
"error": f"HTTP Error: {exc.response.status_code}",
|
||
**credential_context,
|
||
}
|
||
except httpx.HTTPError as exc:
|
||
return {
|
||
"success": False,
|
||
"checksum": checksum,
|
||
"stage": "network",
|
||
"message": f"连接验证失败：{exc.__class__.__name__}",
|
||
"error": str(exc),
|
||
**credential_context,
|
||
}
|
||
|
||
|
||
def make_success_validation(checksum: str, result: dict[str, Any]) -> dict[str, Any]:
|
||
return {
|
||
"checksum": checksum,
|
||
"status": "success",
|
||
"validated_at": datetime.now(UTC).isoformat(),
|
||
"status_code": result.get("status_code"),
|
||
"credential_source": result.get("credential_source"),
|
||
}
|
||
|
||
|
||
def is_builtin_validation_current(config: DataSourceConfig | None, checksum: str) -> bool:
|
||
validation = get_connectivity_validation(config)
|
||
return bool(
|
||
validation
|
||
and validation.get("status") == "success"
|
||
and validation.get("checksum") == checksum
|
||
)
|
||
|
||
|
||
async def get_connectivity_store(db) -> dict[str, Any]:
|
||
result = await db.execute(
|
||
select(SystemSetting).where(SystemSetting.category == CONNECTIVITY_STORE_CATEGORY)
|
||
)
|
||
record = result.scalar_one_or_none()
|
||
return dict(record.payload or {}) if record and isinstance(record.payload, dict) else {}
|
||
|
||
|
||
async def save_connectivity_success(
|
||
db,
|
||
source: str,
|
||
checksum: str,
|
||
result: dict[str, Any],
|
||
*,
|
||
connected_by: str,
|
||
) -> dict[str, Any]:
|
||
store = await get_connectivity_store(db)
|
||
validation = {
|
||
**make_success_validation(checksum, result),
|
||
"connected_by": connected_by,
|
||
}
|
||
store[source] = validation
|
||
|
||
existing = await db.execute(
|
||
select(SystemSetting).where(SystemSetting.category == CONNECTIVITY_STORE_CATEGORY)
|
||
)
|
||
record = existing.scalar_one_or_none()
|
||
if record is None:
|
||
db.add(SystemSetting(category=CONNECTIVITY_STORE_CATEGORY, payload=store))
|
||
else:
|
||
record.payload = store
|
||
return validation
|
||
|
||
|
||
async def load_builtin_override_config(db, source: str) -> DataSourceConfig | None:
|
||
result = await db.execute(
|
||
select(DataSourceConfig)
|
||
.where(DataSourceConfig.name == source)
|
||
.where(DataSourceConfig.is_active.is_(True))
|
||
)
|
||
return result.scalar_one_or_none()
|
||
|
||
|
||
async def get_builtin_effective_candidate(db, source: str) -> dict[str, Any]:
|
||
override = await load_builtin_override_config(db, source)
|
||
default_endpoint = get_data_sources_config().get_yaml_url(source)
|
||
return {
|
||
"name": source,
|
||
"endpoint": (override.endpoint if override and override.endpoint else default_endpoint) or "",
|
||
"auth_type": override.auth_type if override else "none",
|
||
"headers": override.headers if override else {},
|
||
"config": strip_connectivity_validation(override.config if override else {}),
|
||
}
|
||
|
||
|
||
async def has_collected_data(db, source: str) -> bool:
|
||
result = await db.execute(select(func.count(CollectedData.id)).where(CollectedData.source == source))
|
||
if (result.scalar() or 0) > 0:
|
||
return True
|
||
|
||
datasource_result = await db.execute(select(DataSource).where(DataSource.source == source))
|
||
datasource = datasource_result.scalar_one_or_none()
|
||
return bool(datasource and datasource.last_status == "success")
|
||
|
||
|
||
async def get_builtin_connection_status(
|
||
db,
|
||
source: str,
|
||
endpoint: str,
|
||
auth_type: str,
|
||
headers: dict | None,
|
||
config: dict | None,
|
||
) -> dict[str, Any]:
|
||
checksum, credential_context = await build_builtin_connectivity_checksum(
|
||
source,
|
||
endpoint,
|
||
auth_type,
|
||
headers,
|
||
config,
|
||
db,
|
||
)
|
||
store = await get_connectivity_store(db)
|
||
validation = store.get(source)
|
||
if isinstance(validation, dict) and validation.get("status") == "success":
|
||
if validation.get("checksum") == checksum:
|
||
return {
|
||
"connected": True,
|
||
"checksum": checksum,
|
||
"connected_by": validation.get("connected_by") or "connection_button",
|
||
"message": "当前配置已完成连接验证。",
|
||
**credential_context,
|
||
}
|
||
|
||
effective = await get_builtin_effective_candidate(db, source)
|
||
effective_checksum, _ = await build_builtin_connectivity_checksum(
|
||
source,
|
||
effective["endpoint"],
|
||
effective["auth_type"],
|
||
effective["headers"],
|
||
effective["config"],
|
||
db,
|
||
)
|
||
if checksum == effective_checksum and await has_collected_data(db, source):
|
||
return {
|
||
"connected": True,
|
||
"checksum": checksum,
|
||
"connected_by": "collection",
|
||
"message": "当前配置已有成功采集数据，视为已连接。",
|
||
**credential_context,
|
||
}
|
||
|
||
if isinstance(validation, dict) and validation.get("status") == "success":
|
||
return {
|
||
"connected": False,
|
||
"checksum": checksum,
|
||
"connected_by": None,
|
||
"message": "接口地址或凭证指纹已变化，请重新点击连接验证。",
|
||
**credential_context,
|
||
}
|
||
|
||
return {
|
||
"connected": False,
|
||
"checksum": checksum,
|
||
"connected_by": None,
|
||
"message": "当前配置尚未连接，请点击连接验证。",
|
||
**credential_context,
|
||
}
|