Files
sub-provider/app/services/config_store.py
2026-04-20 14:16:09 +08:00

641 lines
22 KiB
Python

from __future__ import annotations
import json
from sqlalchemy import delete, select
from sqlalchemy.orm import Session, selectinload
from app.db import get_session, init_db
from app.db_models import (
AppSettingORM,
PolicyGroupORM,
ProfileORM,
ProfileRuleModuleORM,
ProfileSourceORM,
RuleModuleORM,
SourceORM,
)
from app.models import AppConfig, ClientConfig, ProxyGroupConfig, RegionConfig, RuleConfig, SourceConfig
from app.services.loader import load_yaml_app_config
def _json_dumps(value) -> str:
return json.dumps(value, ensure_ascii=False)
def _json_loads(value: str | None, default):
if not value:
return default
return json.loads(value)
def import_yaml_config_to_db(config_path, *, replace_existing: bool = False) -> None:
app_config = load_yaml_app_config(config_path)
init_db()
with get_session() as session:
if replace_existing:
_clear_config_tables(session)
_upsert_setting(session, "public_path", app_config.public_path or "")
source_rows: list[SourceORM] = []
for key, source in app_config.sources.items():
row = _get_or_create_source(session, key)
row.enabled = source.enabled
row.kind = source.kind
row.url = source.url
row.display_name = source.display_name
row.headers_json = _json_dumps(source.headers)
row.include_regex = source.include_regex
row.exclude_regex = source.exclude_regex
row.prefix = source.prefix
row.suffix = source.suffix
row.cache_ttl_seconds = source.cache_ttl_seconds
source_rows.append(row)
module_rows: list[RuleModuleORM] = []
ordered_rule_keys = list(app_config.rules.keys())
for key, rule in app_config.rules.items():
row = _get_or_create_rule_module(session, key)
row.file_path = rule.file
row.behavior = rule.behavior
row.format = rule.format
row.policy = rule.policy
row.no_resolve = rule.no_resolve
row.payload_json = _json_dumps(rule.payload)
module_rows.append(row)
session.flush()
for client_key, client in app_config.clients.items():
profile = _get_or_create_profile(session, client_key)
_populate_profile(profile, client)
session.flush()
_replace_profile_sources(session, profile.id, source_rows)
_replace_profile_groups(session, profile.id, app_config)
_replace_profile_rule_links(session, profile.id, module_rows, ordered_rule_keys)
session.commit()
def load_app_config_from_db() -> AppConfig | None:
init_db()
with get_session() as session:
return _load_app_config_from_session(session, profile_key=None)
def load_profile_app_config_from_db(profile_key: str) -> AppConfig | None:
init_db()
with get_session() as session:
return _load_app_config_from_session(session, profile_key=profile_key)
def list_sources() -> list[SourceORM]:
init_db()
with get_session() as session:
return list(session.scalars(select(SourceORM).order_by(SourceORM.key)))
def save_source(
*,
key: str,
enabled: bool,
kind: str,
url: str,
display_name: str | None,
headers: dict[str, str],
include_regex: str | None,
exclude_regex: str | None,
prefix: str,
suffix: str,
cache_ttl_seconds: int | None,
) -> None:
init_db()
with get_session() as session:
source = _get_or_create_source(session, key)
source.enabled = enabled
source.kind = kind
source.url = url
source.display_name = display_name
source.headers_json = _json_dumps(headers)
source.include_regex = include_regex
source.exclude_regex = exclude_regex
source.prefix = prefix
source.suffix = suffix
source.cache_ttl_seconds = cache_ttl_seconds
session.commit()
def delete_source(key: str) -> None:
init_db()
with get_session() as session:
source = session.scalar(select(SourceORM).where(SourceORM.key == key))
if source is None:
return
session.execute(delete(ProfileSourceORM).where(ProfileSourceORM.source_id == source.id))
session.delete(source)
session.commit()
def list_profiles() -> list[ProfileORM]:
init_db()
with get_session() as session:
return list(session.scalars(select(ProfileORM).order_by(ProfileORM.key)))
def save_profile(
*,
key: str,
title: str,
provider_interval: int,
rule_interval: int,
test_url: str,
test_interval: int,
main_policy: str,
source_policy: str,
mixed_auto_policy: str,
manual_policy: str,
direct_policy: str,
mode: str,
allow_lan: bool,
ipv6: bool,
mixed_port: int | None,
socks_port: int | None,
log_level: str | None,
) -> None:
init_db()
with get_session() as session:
profile = _get_or_create_profile(session, key)
profile.title = title
profile.provider_interval = provider_interval
profile.rule_interval = rule_interval
profile.test_url = test_url
profile.test_interval = test_interval
profile.main_policy = main_policy
profile.source_policy = source_policy
profile.mixed_auto_policy = mixed_auto_policy
profile.manual_policy = manual_policy
profile.direct_policy = direct_policy
profile.mode = mode
profile.allow_lan = allow_lan
profile.ipv6 = ipv6
profile.mixed_port = mixed_port
profile.socks_port = socks_port
profile.log_level = log_level
session.commit()
def list_profile_rule_bindings(profile_key: str) -> list[dict]:
init_db()
with get_session() as session:
profile = _load_profile_with_relationships(session, profile_key)
if profile is None:
return []
bindings: list[dict] = []
for link in sorted(profile.rule_links, key=lambda item: (item.order_index, item.id)):
module = link.rule_module
bindings.append(
{
"key": module.key,
"enabled": link.enabled,
"order_index": link.order_index,
"policy": link.policy_override or module.policy,
"file_path": module.file_path,
"behavior": module.behavior,
"format": module.format,
}
)
return bindings
def list_profile_source_bindings(profile_key: str) -> list[dict]:
init_db()
with get_session() as session:
profile = _load_profile_with_relationships(session, profile_key)
if profile is None:
return []
bindings: list[dict] = []
for link in sorted(profile.source_links, key=lambda item: (item.order_index, item.id)):
bindings.append(
{
"key": link.source.key,
"enabled": link.enabled,
"order_index": link.order_index,
"display_name": link.source.display_name or link.source.key,
"kind": link.source.kind,
"url": link.source.url,
}
)
return bindings
def update_profile_source_bindings(profile_key: str, rows: list[dict]) -> None:
init_db()
with get_session() as session:
profile = _load_profile_with_relationships(session, profile_key)
if profile is None:
raise KeyError(f"profile not found: {profile_key}")
source_keys = [row["key"] for row in rows]
sources = list(session.scalars(select(SourceORM).where(SourceORM.key.in_(source_keys))))
source_by_key = {source.key: source for source in sources}
session.execute(delete(ProfileSourceORM).where(ProfileSourceORM.profile_id == profile.id))
for row in rows:
source = source_by_key.get(row["key"])
if source is None:
continue
session.add(
ProfileSourceORM(
profile_id=profile.id,
source_id=source.id,
order_index=int(row["order_index"]),
enabled=bool(row["enabled"]),
)
)
session.commit()
def get_profile_default_source_keys(profile_key: str) -> list[str]:
init_db()
with get_session() as session:
profile = _load_profile_with_relationships(session, profile_key)
if profile is None:
return []
keys = [
link.source.key
for link in sorted(profile.source_links, key=lambda item: (item.order_index, item.id))
if link.enabled and link.source.enabled and str(link.source.url).strip()
]
return keys
def list_profile_groups(profile_key: str) -> list[dict]:
init_db()
with get_session() as session:
profile = _load_profile_with_relationships(session, profile_key)
if profile is None:
return []
rows: list[dict] = []
for group in sorted(profile.groups, key=lambda item: (item.order_index, item.id)):
rows.append(
{
"id": group.id,
"group_kind": group.group_kind,
"name": group.name,
"type": group.type,
"order_index": group.order_index,
"proxies": _json_loads(group.proxies_json, []),
"filter_regex": group.filter_regex or "",
"tolerance": group.tolerance,
"url": group.url or "",
"interval": group.interval,
"enabled": group.enabled,
}
)
return rows
def replace_profile_groups(profile_key: str, rows: list[dict]) -> None:
init_db()
with get_session() as session:
profile = _load_profile_with_relationships(session, profile_key)
if profile is None:
raise KeyError(f"profile not found: {profile_key}")
session.execute(delete(PolicyGroupORM).where(PolicyGroupORM.profile_id == profile.id))
for row in rows:
session.add(
PolicyGroupORM(
profile_id=profile.id,
group_kind=row["group_kind"],
name=row["name"],
type=row["type"],
order_index=int(row["order_index"]),
proxies_json=_json_dumps(row["proxies"]) if row["proxies"] else None,
filter_regex=row["filter_regex"] or None,
tolerance=row["tolerance"],
url=row["url"] or None,
interval=row["interval"],
enabled=bool(row["enabled"]),
)
)
session.commit()
def update_profile_rule_bindings(profile_key: str, rows: list[dict]) -> None:
init_db()
with get_session() as session:
profile = _load_profile_with_relationships(session, profile_key)
if profile is None:
raise KeyError(f"profile not found: {profile_key}")
module_by_key = {link.rule_module.key: link for link in profile.rule_links}
for row in rows:
link = module_by_key.get(row["key"])
if link is None:
continue
link.enabled = bool(row["enabled"])
link.order_index = int(row["order_index"])
policy = str(row["policy"]).strip()
link.policy_override = policy or None
session.commit()
def _load_app_config_from_session(session: Session, profile_key: str | None) -> AppConfig | None:
profiles = list(
session.scalars(
select(ProfileORM)
.options(
selectinload(ProfileORM.groups),
selectinload(ProfileORM.rule_links).selectinload(ProfileRuleModuleORM.rule_module),
)
.order_by(ProfileORM.key)
)
)
sources = list(session.scalars(select(SourceORM).order_by(SourceORM.key)))
if not profiles or not sources:
return None
target_profile = next((profile for profile in profiles if profile.key == profile_key), profiles[0])
public_path = _setting_value(session, "public_path")
regions, selector_groups, policy_groups = _build_group_configs(target_profile)
rules = _build_rule_configs(target_profile)
return AppConfig(
public_path=public_path or None,
sources={source.key: _source_to_config(source) for source in sources},
rules=rules,
clients={profile.key: _profile_to_client_config(profile) for profile in profiles},
regions=regions,
selector_groups=selector_groups,
policy_groups=policy_groups,
)
def _source_to_config(source: SourceORM) -> SourceConfig:
return SourceConfig(
enabled=source.enabled,
kind=source.kind,
url=source.url,
display_name=source.display_name,
headers=_json_loads(source.headers_json, {}),
include_regex=source.include_regex,
exclude_regex=source.exclude_regex,
prefix=source.prefix,
suffix=source.suffix,
cache_ttl_seconds=source.cache_ttl_seconds,
)
def _build_group_configs(profile: ProfileORM) -> tuple[dict[str, RegionConfig], list[ProxyGroupConfig], list[ProxyGroupConfig]]:
regions: dict[str, RegionConfig] = {}
selector_groups: list[ProxyGroupConfig] = []
policy_groups: list[ProxyGroupConfig] = []
for group in sorted(profile.groups, key=lambda item: (item.order_index, item.id)):
if not group.enabled:
continue
if group.group_kind == "region":
regions[_slugify(group.name)] = RegionConfig(
name=group.name,
filter=group.filter_regex or "",
tolerance=group.tolerance or 50,
)
continue
target = selector_groups if group.group_kind == "selector" else policy_groups
target.append(
ProxyGroupConfig(
name=group.name,
type=group.type,
proxies=_json_loads(group.proxies_json, []),
filter=group.filter_regex,
tolerance=group.tolerance,
url=group.url,
interval=group.interval,
)
)
return regions, selector_groups, policy_groups
def _build_rule_configs(profile: ProfileORM) -> dict[str, RuleConfig]:
rules: dict[str, RuleConfig] = {}
for link in sorted(profile.rule_links, key=lambda item: (item.order_index, item.id)):
if not link.enabled:
continue
module = link.rule_module
rules[module.key] = RuleConfig(
file=module.file_path,
behavior=module.behavior,
format=module.format,
policy=link.policy_override or module.policy,
no_resolve=module.no_resolve if link.no_resolve_override is None else link.no_resolve_override,
payload=_json_loads(link.payload_override_json, _json_loads(module.payload_json, [])),
)
return rules
def _load_profile_with_relationships(session: Session, profile_key: str) -> ProfileORM | None:
return session.scalar(
select(ProfileORM)
.where(ProfileORM.key == profile_key)
.options(
selectinload(ProfileORM.groups),
selectinload(ProfileORM.rule_links).selectinload(ProfileRuleModuleORM.rule_module),
selectinload(ProfileORM.source_links).selectinload(ProfileSourceORM.source),
)
)
def _clear_config_tables(session: Session) -> None:
session.execute(delete(ProfileRuleModuleORM))
session.execute(delete(PolicyGroupORM))
session.execute(delete(ProfileSourceORM))
session.execute(delete(ProfileORM))
session.execute(delete(RuleModuleORM))
session.execute(delete(SourceORM))
session.execute(delete(AppSettingORM))
session.flush()
def _upsert_setting(session: Session, key: str, value: str) -> None:
row = session.scalar(select(AppSettingORM).where(AppSettingORM.key == key))
if row is None:
row = AppSettingORM(key=key, value=value)
session.add(row)
else:
row.value = value
def _setting_value(session: Session, key: str) -> str | None:
row = session.scalar(select(AppSettingORM).where(AppSettingORM.key == key))
if row is None:
return None
return row.value
def _get_or_create_source(session: Session, key: str) -> SourceORM:
row = session.scalar(select(SourceORM).where(SourceORM.key == key))
if row is None:
row = SourceORM(key=key, kind="auto", url="")
session.add(row)
return row
def _get_or_create_rule_module(session: Session, key: str) -> RuleModuleORM:
row = session.scalar(select(RuleModuleORM).where(RuleModuleORM.key == key))
if row is None:
row = RuleModuleORM(key=key, behavior="classical", format="text", policy="DIRECT")
session.add(row)
return row
def _get_or_create_profile(session: Session, key: str) -> ProfileORM:
row = session.scalar(select(ProfileORM).where(ProfileORM.key == key))
if row is None:
row = ProfileORM(
key=key,
title=key,
test_url="https://www.gstatic.com/generate_204",
main_policy="🚀 节点选择",
source_policy="☁️ 机场选择",
mixed_auto_policy="♻️ 自动选择",
manual_policy="🚀 手动切换",
direct_policy="DIRECT",
)
session.add(row)
return row
def _populate_profile(profile: ProfileORM, client: ClientConfig) -> None:
profile.title = client.title
profile.provider_interval = client.provider_interval
profile.rule_interval = client.rule_interval
profile.test_url = str(client.test_url)
profile.test_interval = client.test_interval
profile.main_policy = client.main_policy
profile.source_policy = client.source_policy
profile.mixed_auto_policy = client.mixed_auto_policy
profile.manual_policy = client.manual_policy
profile.direct_policy = client.direct_policy
profile.mode = client.mode
profile.allow_lan = client.allow_lan
profile.ipv6 = client.ipv6
profile.mixed_port = client.mixed_port
profile.socks_port = client.socks_port
profile.log_level = client.log_level
def _replace_profile_sources(session: Session, profile_id: int, source_rows: list[SourceORM]) -> None:
session.execute(delete(ProfileSourceORM).where(ProfileSourceORM.profile_id == profile_id))
for order_index, source in enumerate(source_rows):
session.add(
ProfileSourceORM(
profile_id=profile_id,
source_id=source.id,
order_index=order_index,
enabled=True,
)
)
def _replace_profile_groups(session: Session, profile_id: int, app_config: AppConfig) -> None:
session.execute(delete(PolicyGroupORM).where(PolicyGroupORM.profile_id == profile_id))
order_index = 0
for region in app_config.regions.values():
session.add(
PolicyGroupORM(
profile_id=profile_id,
group_kind="region",
name=region.name,
type="url-test",
order_index=order_index,
filter_regex=region.filter,
tolerance=region.tolerance,
enabled=True,
)
)
order_index += 1
for group in app_config.selector_groups:
session.add(
PolicyGroupORM(
profile_id=profile_id,
group_kind="selector",
name=group.name,
type=group.type,
order_index=order_index,
proxies_json=_json_dumps(group.proxies),
filter_regex=group.filter,
tolerance=group.tolerance,
url=str(group.url) if group.url else None,
interval=group.interval,
enabled=True,
)
)
order_index += 1
for group in app_config.policy_groups:
session.add(
PolicyGroupORM(
profile_id=profile_id,
group_kind="policy",
name=group.name,
type=group.type,
order_index=order_index,
proxies_json=_json_dumps(group.proxies),
filter_regex=group.filter,
tolerance=group.tolerance,
url=str(group.url) if group.url else None,
interval=group.interval,
enabled=True,
)
)
order_index += 1
def _replace_profile_rule_links(
session: Session,
profile_id: int,
module_rows: list[RuleModuleORM],
ordered_rule_keys: list[str],
) -> None:
session.execute(delete(ProfileRuleModuleORM).where(ProfileRuleModuleORM.profile_id == profile_id))
module_by_key = {module.key: module for module in module_rows}
for order_index, key in enumerate(ordered_rule_keys):
module = module_by_key[key]
session.add(
ProfileRuleModuleORM(
profile_id=profile_id,
rule_module_id=module.id,
enabled=True,
order_index=order_index,
)
)
def _profile_to_client_config(profile: ProfileORM) -> ClientConfig:
return ClientConfig(
title=profile.title,
provider_interval=profile.provider_interval,
rule_interval=profile.rule_interval,
test_url=profile.test_url,
test_interval=profile.test_interval,
main_policy=profile.main_policy,
source_policy=profile.source_policy,
mixed_auto_policy=profile.mixed_auto_policy,
manual_policy=profile.manual_policy,
direct_policy=profile.direct_policy,
mode=profile.mode,
allow_lan=profile.allow_lan,
ipv6=profile.ipv6,
mixed_port=profile.mixed_port,
socks_port=profile.socks_port,
log_level=profile.log_level,
)
def _slugify(name: str) -> str:
text = "".join(char.lower() if char.isalnum() else "-" for char in name).strip("-")
while "--" in text:
text = text.replace("--", "-")
return text or "group"