added py server encryption
This commit is contained in:
parent
a0489a63bb
commit
141f8714c3
@ -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 {}
|
||||
|
||||
@ -478,6 +478,7 @@ func trainingPredictFramePathsBatchForAnalyze(
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
addAIServerAuth(req)
|
||||
|
||||
client := &http.Client{
|
||||
Timeout: 120 * time.Second,
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user