210 lines
7.3 KiB
Python
210 lines
7.3 KiB
Python
"""BarentsWatch AIS credential resolution and connectivity checks."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import shlex
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.core.data_sources import get_data_sources_config
|
||
from app.models.datasource_config import DataSourceConfig
|
||
|
||
|
||
BARENTSWATCH_LATEST_URL = "https://live.ais.barentswatch.no/v1/latest/combined"
|
||
BARENTSWATCH_TOKEN_URL = "https://id.barentswatch.no/connect/token"
|
||
BARENTSWATCH_DATASOURCE_NAME = "barentswatch_vessels"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class BarentsWatchConfig:
|
||
endpoint: str
|
||
client_id: str
|
||
client_secret: str
|
||
credential_source: str
|
||
endpoint_source: str
|
||
|
||
|
||
def _read_zshrc_env(path: Path | None = None) -> dict[str, str]:
|
||
zshrc_path = path or Path.home() / ".zshrc"
|
||
if not zshrc_path.exists():
|
||
return {}
|
||
|
||
values: dict[str, str] = {}
|
||
for raw_line in zshrc_path.read_text(encoding="utf-8", errors="ignore").splitlines():
|
||
line = raw_line.strip()
|
||
if not line or line.startswith("#"):
|
||
continue
|
||
if line.startswith("export "):
|
||
line = line[len("export ") :].strip()
|
||
if "=" not in line:
|
||
continue
|
||
|
||
key, value = line.split("=", 1)
|
||
key = key.strip()
|
||
if not key or not key.replace("_", "").isalnum() or not key[0].isalpha():
|
||
continue
|
||
|
||
try:
|
||
parsed = shlex.split(value, comments=True, posix=True)
|
||
except ValueError:
|
||
parsed = [value.strip().strip("'\"")]
|
||
if parsed:
|
||
values[key] = parsed[0]
|
||
return values
|
||
|
||
|
||
def _first_env_value(zshrc_env: dict[str, str], *keys: str) -> tuple[str, str]:
|
||
for key in keys:
|
||
value = os.getenv(key)
|
||
if value:
|
||
return value, "environment"
|
||
for key in keys:
|
||
value = zshrc_env.get(key)
|
||
if value:
|
||
return value, "~/.zshrc"
|
||
return "", ""
|
||
|
||
|
||
async def get_barentswatch_datasource_record(db: AsyncSession) -> DataSourceConfig | None:
|
||
result = await db.execute(
|
||
select(DataSourceConfig)
|
||
.where(DataSourceConfig.name == BARENTSWATCH_DATASOURCE_NAME)
|
||
.where(DataSourceConfig.is_active.is_(True))
|
||
)
|
||
return result.scalar_one_or_none()
|
||
|
||
|
||
async def resolve_barentswatch_config(db: AsyncSession | None = None) -> BarentsWatchConfig:
|
||
record = await get_barentswatch_datasource_record(db) if db else None
|
||
auth_config = dict(record.auth_config or {}) if record else {}
|
||
config = dict(record.config or {}) if record else {}
|
||
zshrc_env = _read_zshrc_env()
|
||
|
||
env_client_id, env_source = _first_env_value(
|
||
zshrc_env,
|
||
"BARENTSWATCH_CLIENT_ID",
|
||
"BARRENTSWATCH_CLIENT_ID",
|
||
)
|
||
env_client_secret, secret_env_source = _first_env_value(
|
||
zshrc_env,
|
||
"BARENTSWATCH_CLIENT_SECRET",
|
||
"BARRENTSWATCH_CLIENT_SECRET",
|
||
)
|
||
client_id = auth_config.get("client_id") or config.get("client_id") or env_client_id
|
||
client_secret = (
|
||
auth_config.get("client_secret") or config.get("client_secret") or env_client_secret
|
||
)
|
||
|
||
credential_source = ""
|
||
if auth_config.get("client_id") or auth_config.get("client_secret"):
|
||
credential_source = "datasource_config"
|
||
elif config.get("client_id") or config.get("client_secret"):
|
||
credential_source = "datasource_runtime_config"
|
||
elif env_source or secret_env_source:
|
||
credential_source = env_source or secret_env_source
|
||
|
||
yaml_endpoint = get_data_sources_config().get_yaml_url(BARENTSWATCH_DATASOURCE_NAME)
|
||
endpoint = record.endpoint if record and record.endpoint else yaml_endpoint
|
||
return BarentsWatchConfig(
|
||
endpoint=endpoint or BARENTSWATCH_LATEST_URL,
|
||
client_id=str(client_id or ""),
|
||
client_secret=str(client_secret or ""),
|
||
credential_source=credential_source or "missing",
|
||
endpoint_source="datasource_config" if record and record.endpoint else "default",
|
||
)
|
||
|
||
|
||
async def fetch_barentswatch_access_token(
|
||
client: httpx.AsyncClient,
|
||
config: BarentsWatchConfig,
|
||
) -> str | None:
|
||
if not config.client_id or not config.client_secret:
|
||
return None
|
||
|
||
response = await client.post(
|
||
BARENTSWATCH_TOKEN_URL,
|
||
data={
|
||
"client_id": config.client_id,
|
||
"client_secret": config.client_secret,
|
||
"scope": "ais",
|
||
"grant_type": "client_credentials",
|
||
},
|
||
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
||
)
|
||
response.raise_for_status()
|
||
payload = response.json()
|
||
token = payload.get("access_token")
|
||
return str(token) if token else None
|
||
|
||
|
||
async def check_barentswatch_connectivity(db: AsyncSession) -> dict[str, Any]:
|
||
config = await resolve_barentswatch_config(db)
|
||
return await check_barentswatch_config(config)
|
||
|
||
|
||
async def check_barentswatch_config(config: BarentsWatchConfig) -> dict[str, Any]:
|
||
if not config.client_id or not config.client_secret:
|
||
return {
|
||
"success": False,
|
||
"stage": "credentials",
|
||
"message": "未找到 BarentsWatch client id/client secret，请先配置采集器凭证。",
|
||
"endpoint": config.endpoint,
|
||
"credential_source": config.credential_source,
|
||
"settings_tab": "collector_credentials",
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=20.0) as client:
|
||
token = await fetch_barentswatch_access_token(client, config)
|
||
if not token:
|
||
return {
|
||
"success": False,
|
||
"stage": "token",
|
||
"message": "BarentsWatch token 响应中没有 access_token，请检查凭证。",
|
||
"endpoint": config.endpoint,
|
||
"credential_source": config.credential_source,
|
||
"settings_tab": "collector_credentials",
|
||
}
|
||
|
||
async with client.stream(
|
||
"GET",
|
||
config.endpoint,
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
) as response:
|
||
response.raise_for_status()
|
||
|
||
return {
|
||
"success": True,
|
||
"stage": "endpoint",
|
||
"message": "BarentsWatch AIS token 和数据接口均可连通。",
|
||
"endpoint": config.endpoint,
|
||
"credential_source": config.credential_source,
|
||
"endpoint_source": config.endpoint_source,
|
||
}
|
||
except httpx.HTTPStatusError as exc:
|
||
status_code = exc.response.status_code
|
||
stage = "token" if str(exc.request.url) == BARENTSWATCH_TOKEN_URL else "endpoint"
|
||
return {
|
||
"success": False,
|
||
"stage": stage,
|
||
"message": f"BarentsWatch {stage} 请求返回 HTTP {status_code}，请检查凭证或接口地址。",
|
||
"endpoint": config.endpoint,
|
||
"credential_source": config.credential_source,
|
||
"settings_tab": "collector_credentials",
|
||
}
|
||
except httpx.HTTPError as exc:
|
||
return {
|
||
"success": False,
|
||
"stage": "network",
|
||
"message": f"BarentsWatch 链路检查失败：{exc.__class__.__name__}",
|
||
"endpoint": config.endpoint,
|
||
"credential_source": config.credential_source,
|
||
"settings_tab": "collector_credentials",
|
||
}
|