added py server encryption
This commit is contained in:
parent
a0489a63bb
commit
141f8714c3
@ -2,10 +2,11 @@
|
|||||||
|
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
|
import secrets
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import List, Optional
|
from typing import List, Optional
|
||||||
|
|
||||||
from fastapi import FastAPI
|
from fastapi import Depends, FastAPI, HTTPException, Request, status
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from ultralytics import YOLO
|
from ultralytics import YOLO
|
||||||
|
|
||||||
@ -125,6 +126,42 @@ model = None
|
|||||||
|
|
||||||
app = FastAPI()
|
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):
|
class PredictBatchRequest(BaseModel):
|
||||||
paths: List[str]
|
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):
|
def predict_batch(req: PredictBatchRequest):
|
||||||
paths = [str(path).strip() for path in req.paths if str(path).strip()]
|
paths = [str(path).strip() for path in req.paths if str(path).strip()]
|
||||||
if not paths:
|
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():
|
def health():
|
||||||
current_model = get_model()
|
current_model = get_model()
|
||||||
names = getattr(current_model, "names", {}) or {} if current_model is not None else {}
|
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")
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
addAIServerAuth(req)
|
||||||
|
|
||||||
client := &http.Client{
|
client := &http.Client{
|
||||||
Timeout: 120 * time.Second,
|
Timeout: 120 * time.Second,
|
||||||
|
|||||||
@ -4,6 +4,8 @@ package main
|
|||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
|
"crypto/rand"
|
||||||
|
"encoding/hex"
|
||||||
"fmt"
|
"fmt"
|
||||||
"net"
|
"net"
|
||||||
"net/http"
|
"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 {
|
func aiServerURL() string {
|
||||||
raw := strings.TrimSpace(os.Getenv("AI_SERVER_URL"))
|
raw := strings.TrimSpace(os.Getenv("AI_SERVER_URL"))
|
||||||
if raw == "" {
|
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)
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url+"/health", nil)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
addAIServerAuth(req)
|
||||||
|
|
||||||
res, err := client.Do(req)
|
res, err := client.Do(req)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
_ = res.Body.Close()
|
_ = res.Body.Close()
|
||||||
@ -263,6 +314,8 @@ func aiServerStatusHandler(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
addAIServerAuth(req)
|
||||||
|
|
||||||
res, err := client.Do(req)
|
res, err := client.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
status["error"] = err.Error()
|
status["error"] = err.Error()
|
||||||
@ -334,6 +387,11 @@ func startAIServer(ctx context.Context) (*aiServerProcess, error) {
|
|||||||
"AI_SERVER_URL="+aiServerURL(),
|
"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.
|
// Immer App-Root/generated/training verwenden.
|
||||||
env = upsertEnv(env, "APP_BASE_DIR", appBaseDir)
|
env = upsertEnv(env, "APP_BASE_DIR", appBaseDir)
|
||||||
env = upsertEnv(env, "TRAINING_ROOT", trainingRoot)
|
env = upsertEnv(env, "TRAINING_ROOT", trainingRoot)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user