359 lines
12 KiB
Python
359 lines
12 KiB
Python
import uuid
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import List, Optional
|
|
from fastapi import APIRouter, Depends, HTTPException, status, Request
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy import select, func, desc
|
|
|
|
from app.core.database import get_db
|
|
from app.core.security import generate_registration_code, generate_device_token, hash_token
|
|
from app.models.models import User, Client, ClientCredential, RegistrationCode
|
|
from app.schemas.schemas import (
|
|
ClientResponse, ClientRegisterRequest, ClientRegisterResponse,
|
|
RegistrationCodeCreate, RegistrationCodeResponse, ClientHeartbeatRequest,
|
|
ClientUpdateRequest
|
|
)
|
|
from app.api.deps import get_current_user, require_admin, get_current_client
|
|
from app.services.event_service import log_event
|
|
from app.ws.manager import ws_manager
|
|
|
|
router = APIRouter(prefix="/clients", tags=["Clients"])
|
|
|
|
@router.get("", response_model=List[ClientResponse])
|
|
async def list_clients(
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: User = Depends(get_current_user)
|
|
):
|
|
result = await db.execute(select(Client).order_by(desc(Client.created_at)))
|
|
return result.scalars().all()
|
|
|
|
@router.post("/registration-code", response_model=RegistrationCodeResponse)
|
|
async def create_registration_code(
|
|
payload: RegistrationCodeCreate,
|
|
db: AsyncSession = Depends(get_db),
|
|
admin_user: User = Depends(require_admin)
|
|
):
|
|
code_str = generate_registration_code()
|
|
expires = datetime.now(timezone.utc) + timedelta(hours=payload.expires_in_hours)
|
|
|
|
reg_code = RegistrationCode(
|
|
code=code_str,
|
|
client_name_hint=payload.client_name_hint,
|
|
expires_at=expires,
|
|
is_used=False
|
|
)
|
|
db.add(reg_code)
|
|
await db.commit()
|
|
await db.refresh(reg_code)
|
|
|
|
await log_event(
|
|
db=db,
|
|
event_type="REGISTRATION_CODE_GENERATED",
|
|
message=f"Generated registration code {code_str} (hint: {payload.client_name_hint or 'None'}).",
|
|
severity="INFO",
|
|
user_email=admin_user.email
|
|
)
|
|
|
|
return reg_code
|
|
|
|
@router.post("/register", response_model=ClientRegisterResponse)
|
|
async def register_client(
|
|
payload: ClientRegisterRequest,
|
|
request: Request,
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""Called by the Windows Agent during initial setup to register against the server."""
|
|
# Find valid registration code
|
|
now = datetime.now(timezone.utc)
|
|
res = await db.execute(
|
|
select(RegistrationCode).where(
|
|
RegistrationCode.code == payload.registration_code.strip().upper(),
|
|
RegistrationCode.is_used == False,
|
|
RegistrationCode.expires_at > now
|
|
)
|
|
)
|
|
reg_code = res.scalar_one_or_none()
|
|
if not reg_code:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="Invalid, expired, or already used registration code."
|
|
)
|
|
|
|
client_ip = request.client.host if request.client else None
|
|
|
|
if reg_code.client_id:
|
|
# Re-registering an existing client!
|
|
client_res = await db.execute(select(Client).where(Client.id == reg_code.client_id))
|
|
client = client_res.scalar_one_or_none()
|
|
if not client:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Client associated with this registration code not found."
|
|
)
|
|
# Update connection info and mark as active
|
|
client.hostname = payload.hostname
|
|
client.os_info = payload.os_info
|
|
client.ip_address = client_ip
|
|
client.agent_version = payload.agent_version
|
|
client.status = "ONLINE"
|
|
client.last_seen_at = now
|
|
client.is_active = True
|
|
else:
|
|
# Create a new Client
|
|
max_id_res = await db.execute(select(func.max(Client.id)))
|
|
max_id = max_id_res.scalar() or 0
|
|
client_code = f"CLIENT-{max_id + 1:04d}"
|
|
|
|
# Load default client quota from settings
|
|
from app.services.settings_service import get_setting
|
|
default_quota_gb_str = await get_setting(db, "default_client_quota_gb")
|
|
default_quota_bytes = 100 * 1024 * 1024 * 1024
|
|
if default_quota_gb_str:
|
|
try:
|
|
default_quota_bytes = int(default_quota_gb_str) * 1024 * 1024 * 1024
|
|
except ValueError:
|
|
pass
|
|
|
|
client = Client(
|
|
client_code=client_code,
|
|
name=payload.name or reg_code.client_name_hint or payload.hostname,
|
|
hostname=payload.hostname,
|
|
os_info=payload.os_info,
|
|
ip_address=client_ip,
|
|
agent_version=payload.agent_version,
|
|
status="ONLINE",
|
|
last_seen_at=now,
|
|
is_active=True,
|
|
storage_quota_bytes=default_quota_bytes
|
|
)
|
|
db.add(client)
|
|
await db.flush()
|
|
|
|
# Generate device unique ID and secret token
|
|
device_id = str(uuid.uuid4())
|
|
device_token = generate_device_token()
|
|
token_hash = hash_token(device_token)
|
|
|
|
credential = ClientCredential(
|
|
client_id=client.id,
|
|
device_id=device_id,
|
|
token_hash=token_hash,
|
|
name=f"{payload.hostname} Agent",
|
|
is_revoked=False,
|
|
created_at=now,
|
|
last_used_at=now
|
|
)
|
|
db.add(credential)
|
|
|
|
# Mark registration code as used
|
|
reg_code.is_used = True
|
|
await db.commit()
|
|
|
|
await log_event(
|
|
db=db,
|
|
event_type="CLIENT_REGISTERED",
|
|
message=f"Windows client registered: {client.name} ({client.client_code}, Hostname: {client.hostname}, IP: {client_ip})",
|
|
severity="INFO",
|
|
client_id=client.id,
|
|
ip_address=client_ip,
|
|
details={"device_id": device_id, "os_info": payload.os_info}
|
|
)
|
|
|
|
await ws_manager.broadcast("CLIENT_REGISTERED", {
|
|
"id": client.id,
|
|
"client_code": client.client_code,
|
|
"name": client.name,
|
|
"hostname": client.hostname,
|
|
"status": client.status
|
|
})
|
|
|
|
return {
|
|
"client_code": client.client_code,
|
|
"device_id": device_id,
|
|
"device_token": device_token,
|
|
"name": client.name,
|
|
"server_time": now
|
|
}
|
|
|
|
@router.post("/{client_id}/heartbeat")
|
|
async def client_heartbeat(
|
|
client_id: int,
|
|
payload: ClientHeartbeatRequest,
|
|
request: Request,
|
|
db: AsyncSession = Depends(get_db),
|
|
current_client: Client = Depends(get_current_client)
|
|
):
|
|
"""Heartbeat endpoint invoked periodically by the Windows Agent."""
|
|
if current_client.id != client_id:
|
|
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Client ID mismatch")
|
|
|
|
now = datetime.now(timezone.utc)
|
|
current_client.last_seen_at = now
|
|
current_client.status = payload.status
|
|
if payload.agent_version:
|
|
current_client.agent_version = payload.agent_version
|
|
|
|
client_ip = payload.ip_address or (request.client.host if request.client else None)
|
|
if client_ip:
|
|
current_client.ip_address = client_ip
|
|
|
|
await db.commit()
|
|
|
|
await ws_manager.broadcast("CLIENT_HEARTBEAT", {
|
|
"client_id": current_client.id,
|
|
"status": current_client.status,
|
|
"last_seen_at": now.isoformat()
|
|
})
|
|
|
|
return {"status": "ok", "server_time": now}
|
|
|
|
@router.get("/{client_id}", response_model=ClientResponse)
|
|
async def get_client(
|
|
client_id: int,
|
|
db: AsyncSession = Depends(get_db),
|
|
current_user: User = Depends(get_current_user)
|
|
):
|
|
res = await db.execute(select(Client).where(Client.id == client_id))
|
|
client = res.scalar_one_or_none()
|
|
if not client:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Client not found")
|
|
return client
|
|
|
|
@router.post("/{client_id}/revoke")
|
|
async def revoke_client_credentials(
|
|
client_id: int,
|
|
db: AsyncSession = Depends(get_db),
|
|
admin_user: User = Depends(require_admin)
|
|
):
|
|
"""Revokes all active authentication tokens for a client."""
|
|
res = await db.execute(select(ClientCredential).where(ClientCredential.client_id == client_id))
|
|
creds = res.scalars().all()
|
|
for cred in creds:
|
|
cred.is_revoked = True
|
|
|
|
client_res = await db.execute(select(Client).where(Client.id == client_id))
|
|
client = client_res.scalar_one_or_none()
|
|
if client:
|
|
client.status = "OFFLINE"
|
|
|
|
await db.commit()
|
|
|
|
await log_event(
|
|
db=db,
|
|
event_type="CREDENTIALS_REVOKED",
|
|
message=f"Revoked credentials for client ID {client_id}.",
|
|
severity="WARNING",
|
|
client_id=client_id,
|
|
user_email=admin_user.email
|
|
)
|
|
|
|
return {"message": "Client credentials successfully revoked"}
|
|
|
|
@router.post("/{client_id}/re-register", response_model=RegistrationCodeResponse)
|
|
async def generate_client_re_registration_code(
|
|
client_id: int,
|
|
db: AsyncSession = Depends(get_db),
|
|
admin_user: User = Depends(require_admin)
|
|
):
|
|
"""Generates a registration code to re-link an existing client's agent."""
|
|
client_res = await db.execute(select(Client).where(Client.id == client_id))
|
|
client = client_res.scalar_one_or_none()
|
|
if not client:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="Client not found"
|
|
)
|
|
|
|
# Revoke old credentials to prepare for the new agent connection
|
|
cred_res = await db.execute(select(ClientCredential).where(ClientCredential.client_id == client_id))
|
|
creds = cred_res.scalars().all()
|
|
for cred in creds:
|
|
cred.is_revoked = True
|
|
|
|
code_str = generate_registration_code()
|
|
expires = datetime.now(timezone.utc) + timedelta(hours=24) # 24 hours to re-register
|
|
|
|
reg_code = RegistrationCode(
|
|
code=code_str,
|
|
client_id=client.id,
|
|
client_name_hint=client.name,
|
|
expires_at=expires,
|
|
is_used=False
|
|
)
|
|
db.add(reg_code)
|
|
await db.commit()
|
|
await db.refresh(reg_code)
|
|
|
|
await log_event(
|
|
db=db,
|
|
event_type="REGISTRATION_CODE_GENERATED",
|
|
message=f"Generated re-registration code {code_str} for client {client.name} ({client.client_code}).",
|
|
severity="INFO",
|
|
client_id=client_id,
|
|
user_email=admin_user.email
|
|
)
|
|
|
|
return reg_code
|
|
|
|
@router.patch("/{client_id}", response_model=ClientResponse)
|
|
async def update_client(
|
|
client_id: int,
|
|
payload: ClientUpdateRequest,
|
|
db: AsyncSession = Depends(get_db),
|
|
admin_user: User = Depends(require_admin)
|
|
):
|
|
"""Updates client details such as alias."""
|
|
res = await db.execute(select(Client).where(Client.id == client_id))
|
|
client = res.scalar_one_or_none()
|
|
if not client:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Client not found")
|
|
|
|
if payload.alias is not None:
|
|
client.alias = payload.alias
|
|
if payload.storage_quota_bytes is not None:
|
|
client.storage_quota_bytes = payload.storage_quota_bytes
|
|
|
|
await db.commit()
|
|
await db.refresh(client)
|
|
|
|
message_parts = []
|
|
if payload.alias is not None:
|
|
message_parts.append(f"alias to '{client.alias}'")
|
|
if payload.storage_quota_bytes is not None:
|
|
message_parts.append(f"quota to {client.storage_quota_bytes} bytes")
|
|
|
|
await log_event(
|
|
db=db,
|
|
event_type="CLIENT_UPDATED",
|
|
message=f"Updated client {client.client_code}: {', '.join(message_parts)}.",
|
|
severity="INFO",
|
|
client_id=client.id,
|
|
user_email=admin_user.email
|
|
)
|
|
|
|
return client
|
|
|
|
@router.delete("/{client_id}")
|
|
async def delete_client(
|
|
client_id: int,
|
|
db: AsyncSession = Depends(get_db),
|
|
admin_user: User = Depends(require_admin)
|
|
):
|
|
res = await db.execute(select(Client).where(Client.id == client_id))
|
|
client = res.scalar_one_or_none()
|
|
if not client:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Client not found")
|
|
|
|
await db.delete(client)
|
|
await db.commit()
|
|
|
|
await log_event(
|
|
db=db,
|
|
event_type="CLIENT_DELETED",
|
|
message=f"Deleted client {client.name} ({client.client_code}).",
|
|
severity="WARNING",
|
|
user_email=admin_user.email
|
|
)
|
|
|
|
return {"message": f"Client {client.client_code} deleted"}
|