added py server encryption

This commit is contained in:
Linrador 2026-05-12 14:54:02 +02:00
parent a0489a63bb
commit 141f8714c3
3 changed files with 99 additions and 3 deletions

View File

@ -2,10 +2,11 @@
import json
import os
import secrets
from pathlib import Path
from typing import List, Optional
from fastapi import FastAPI
from fastapi import Depends, FastAPI, HTTPException, Request, status
from pydantic import BaseModel
from ultralytics import YOLO
@ -125,6 +126,42 @@ model = None
app = FastAPI()
AI_SERVER_TOKEN = os.environ.get("AI_SERVER_TOKEN", "").strip()
AI_SERVER_AUTH_REQUIRED = os.environ.get("AI_SERVER_AUTH_REQUIRED", "1").strip().lower() not in {
"0",
"false",
"no",
"off",
}
def require_ai_server_auth(request: Request) -> None:
if not AI_SERVER_AUTH_REQUIRED:
return
if not AI_SERVER_TOKEN:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="AI server auth token not configured",
)
auth = request.headers.get("authorization", "").strip()
header_token = request.headers.get("x-ai-server-token", "").strip()
provided = ""
if auth.lower().startswith("bearer "):
provided = auth[7:].strip()
elif header_token:
provided = header_token
if not provided or not secrets.compare_digest(provided, AI_SERVER_TOKEN):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="unauthorized",
headers={"WWW-Authenticate": "Bearer"},
)
class PredictBatchRequest(BaseModel):
paths: List[str]
@ -349,7 +386,7 @@ def prediction_from_result(result) -> dict:
}
@app.post("/predict-batch")
@app.post("/predict-batch", dependencies=[Depends(require_ai_server_auth)])
def predict_batch(req: PredictBatchRequest):
paths = [str(path).strip() for path in req.paths if str(path).strip()]
if not paths:
@ -402,7 +439,7 @@ def predict_batch(req: PredictBatchRequest):
}
@app.get("/health")
@app.get("/health", dependencies=[Depends(require_ai_server_auth)])
def health():
current_model = get_model()
names = getattr(current_model, "names", {}) or {} if current_model is not None else {}

View File

@ -478,6 +478,7 @@ func trainingPredictFramePathsBatchForAnalyze(
}
req.Header.Set("Content-Type", "application/json")
addAIServerAuth(req)
client := &http.Client{
Timeout: 120 * time.Second,

View File

@ -4,6 +4,8 @@ package main
import (
"bytes"
"context"
"crypto/rand"
"encoding/hex"
"fmt"
"net"
"net/http"
@ -119,6 +121,53 @@ func aiServerAutostartEnabled() bool {
}
}
var aiServerTokenOnce sync.Once
var aiServerTokenValue string
func randomAISecretHex(bytesLen int) string {
if bytesLen <= 0 {
bytesLen = 32
}
b := make([]byte, bytesLen)
if _, err := rand.Read(b); err != nil {
fallback := fmt.Sprintf("%d-%d", time.Now().UnixNano(), os.Getpid())
return hex.EncodeToString([]byte(fallback))
}
return hex.EncodeToString(b)
}
func aiServerToken() string {
aiServerTokenOnce.Do(func() {
token := strings.TrimSpace(os.Getenv("AI_SERVER_TOKEN"))
if token == "" {
token = randomAISecretHex(32)
_ = os.Setenv("AI_SERVER_TOKEN", token)
appLogln("🔐 AI_SERVER_TOKEN automatisch erzeugt.")
}
aiServerTokenValue = token
})
return aiServerTokenValue
}
func addAIServerAuth(req *http.Request) {
if req == nil {
return
}
token := strings.TrimSpace(aiServerToken())
if token == "" {
return
}
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("X-AI-Server-Token", token)
}
func aiServerURL() string {
raw := strings.TrimSpace(os.Getenv("AI_SERVER_URL"))
if raw == "" {
@ -221,6 +270,8 @@ func waitForAIServer(ctx context.Context, url string, timeout time.Duration) err
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url+"/health", nil)
if err == nil {
addAIServerAuth(req)
res, err := client.Do(req)
if err == nil {
_ = res.Body.Close()
@ -263,6 +314,8 @@ func aiServerStatusHandler(w http.ResponseWriter, r *http.Request) {
return
}
addAIServerAuth(req)
res, err := client.Do(req)
if err != nil {
status["error"] = err.Error()
@ -334,6 +387,11 @@ func startAIServer(ctx context.Context) (*aiServerProcess, error) {
"AI_SERVER_URL="+aiServerURL(),
)
env = upsertEnv(env, "AI_SERVER_AUTH_REQUIRED", "1")
env = upsertEnv(env, "AI_SERVER_TOKEN", aiServerToken())
appLogln("🔐 AI Server Auth aktiv.")
// Immer App-Root/generated/training verwenden.
env = upsertEnv(env, "APP_BASE_DIR", appBaseDir)
env = upsertEnv(env, "TRAINING_ROOT", trainingRoot)