96 lines
3.8 KiB
Python
96 lines
3.8 KiB
Python
"""Private controller configuration. Never expose credentials through GET APIs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import secrets
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
from cryptography.fernet import Fernet
|
|
from pydantic import BaseModel, Field, field_validator
|
|
|
|
|
|
class PoolSettings(BaseModel):
|
|
enabled: bool = False
|
|
transport: Literal["ssh", "socket"] = "ssh"
|
|
host: str = ""
|
|
port: int = Field(default=22, ge=1, le=65535)
|
|
username: str = ""
|
|
host_key_sha256: str = ""
|
|
source_container: str = ""
|
|
network: str = ""
|
|
python_executable: str = Field(default="/opt/venv/bin/python", pattern=r"^/[a-zA-Z0-9_./-]+$")
|
|
working_dir: str = Field(default="/app/backend", pattern=r"^/[a-zA-Z0-9_./-]+$")
|
|
desired: int = Field(default=0, ge=0, le=32)
|
|
concurrency: int = Field(default=4, ge=1, le=32)
|
|
database_pool_size: int = Field(default=4, ge=1, le=32)
|
|
postgres_pool_size: int = Field(default=4, ge=1, le=32)
|
|
database_connection_budget: int = Field(default=64, ge=3, le=1024)
|
|
queue_limit: int = Field(default=128, ge=1, le=1024)
|
|
queue_timeout: int = Field(default=60, ge=1, le=300)
|
|
cpu: float = Field(default=2, ge=0.25, le=64)
|
|
memory_mb: int = Field(default=2048, ge=512, le=131072)
|
|
total_cpu: float = Field(default=8, ge=0.25, le=512)
|
|
total_memory_mb: int = Field(default=8192, ge=512, le=1048576)
|
|
password: str | None = Field(default=None, max_length=4096, repr=False)
|
|
|
|
@field_validator("source_container", "network")
|
|
@classmethod
|
|
def docker_name(cls, value: str) -> str:
|
|
if value and not re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9_.-]{0,127}", value):
|
|
raise ValueError("容器名和网络名只能包含字母、数字、点、横线、下划线")
|
|
return value
|
|
|
|
|
|
class SettingsStore:
|
|
def __init__(self, directory: Path | None = None):
|
|
self.directory = directory or Path(os.getenv("DEER_FLOW_WORKER_CONTROL_HOME", ".worker-pool")).resolve()
|
|
self.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
key_path = self.directory / "secret.key"
|
|
try:
|
|
fd = os.open(key_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
|
except FileExistsError:
|
|
pass
|
|
else:
|
|
with os.fdopen(fd, "wb") as out:
|
|
out.write(Fernet.generate_key())
|
|
self.cipher = Fernet(key_path.read_bytes())
|
|
self.path = self.directory / "settings.json"
|
|
if not self.path.exists():
|
|
self._write({"instance": secrets.token_hex(8), "token": secrets.token_urlsafe(48), "settings": {}})
|
|
|
|
def _write(self, data: dict) -> None:
|
|
temp = self.path.with_suffix(".tmp")
|
|
fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
|
with os.fdopen(fd, "w", encoding="utf-8") as out:
|
|
json.dump(data, out)
|
|
out.flush()
|
|
os.fsync(out.fileno())
|
|
os.replace(temp, self.path)
|
|
|
|
def raw(self) -> dict:
|
|
return json.loads(self.path.read_text(encoding="utf-8"))
|
|
|
|
def load(self) -> PoolSettings:
|
|
data = self.raw()["settings"].copy()
|
|
encrypted = data.pop("encrypted_password", None)
|
|
if encrypted:
|
|
data["password"] = self.cipher.decrypt(encrypted.encode()).decode()
|
|
return PoolSettings(**data)
|
|
|
|
def save(self, settings: PoolSettings) -> None:
|
|
data = self.raw()
|
|
values = settings.model_dump(exclude={"password"})
|
|
password = settings.password if settings.password is not None else self.load().password
|
|
if password:
|
|
values["encrypted_password"] = self.cipher.encrypt(password.encode()).decode()
|
|
data["settings"] = values
|
|
self._write(data)
|
|
|
|
def public(self) -> dict:
|
|
settings = self.load()
|
|
return {**settings.model_dump(exclude={"password"}), "password_set": bool(settings.password)}
|