Remote-Terminal-for-MeshCore/app/routers/packets.py

import logging
from hashlib import sha256

from fastapi import APIRouter, BackgroundTasks
from pydantic import BaseModel, Field

from app.decoder import try_decrypt_packet_with_channel_key
from app.packet_processor import create_message_from_decrypted
from app.repository import RawPacketRepository

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/packets", tags=["packets"])


class DecryptRequest(BaseModel):
    key_type: str = Field(description="Type of key: 'channel' or 'contact'")
    channel_key: str | None = Field(default=None, description="Channel key as hex (16 bytes = 32 chars)")
    channel_name: str | None = Field(default=None, description="Channel name (for hashtag channels, key derived from name)")


class DecryptResult(BaseModel):
    started: bool
    total_packets: int
    message: str


class DecryptProgress(BaseModel):
    total: int
    processed: int
    decrypted: int
    in_progress: bool


# Global state for tracking decryption progress
_decrypt_progress: DecryptProgress | None = None


async def _run_historical_decryption(channel_key_bytes: bytes, channel_key_hex: str) -> None:
    """Background task to decrypt historical packets with a channel key."""
    global _decrypt_progress

    packets = await RawPacketRepository.get_all_undecrypted()
    total = len(packets)
    processed = 0
    decrypted_count = 0

    _decrypt_progress = DecryptProgress(
        total=total, processed=0, decrypted=0, in_progress=True
    )

    logger.info("Starting historical decryption of %d packets", total)

    for packet_id, packet_data, packet_timestamp in packets:
        result = try_decrypt_packet_with_channel_key(packet_data, channel_key_bytes)

        if result is not None:
            # Successfully decrypted - use shared logic to store message
            logger.debug(
                "Decrypted packet %d: sender=%s, message=%s",
                packet_id,
                result.sender,
                result.message[:50] if result.message else "",
            )

            msg_id = await create_message_from_decrypted(
                packet_id=packet_id,
                channel_key=channel_key_hex,
                sender=result.sender,
                message_text=result.message,
                timestamp=result.timestamp,
                received_at=packet_timestamp,  # Use original packet timestamp for correct ordering
            )

            if msg_id is not None:
                decrypted_count += 1

        processed += 1
        _decrypt_progress = DecryptProgress(
            total=total, processed=processed, decrypted=decrypted_count, in_progress=True
        )

    _decrypt_progress = DecryptProgress(
        total=total, processed=processed, decrypted=decrypted_count, in_progress=False
    )

    logger.info(
        "Historical decryption complete: %d/%d packets decrypted", decrypted_count, total
    )


@router.get("/undecrypted/count")
async def get_undecrypted_count() -> dict:
    """Get the count of undecrypted packets."""
    count = await RawPacketRepository.get_undecrypted_count()
    return {"count": count}


@router.post("/decrypt/historical", response_model=DecryptResult)
async def decrypt_historical_packets(
    request: DecryptRequest, background_tasks: BackgroundTasks
) -> DecryptResult:
    """
    Attempt to decrypt historical packets with the provided key.
    Runs in the background to avoid blocking.
    """
    global _decrypt_progress

    # Check if decryption is already in progress
    if _decrypt_progress and _decrypt_progress.in_progress:
        return DecryptResult(
            started=False,
            total_packets=_decrypt_progress.total,
            message=f"Decryption already in progress: {_decrypt_progress.processed}/{_decrypt_progress.total}",
        )

    # Determine the channel key
    channel_key_bytes: bytes | None = None
    channel_key_hex: str | None = None

    if request.key_type == "channel":
        if request.channel_key:
            # Direct key provided
            try:
                channel_key_bytes = bytes.fromhex(request.channel_key)
                if len(channel_key_bytes) != 16:
                    return DecryptResult(
                        started=False,
                        total_packets=0,
                        message="Channel key must be 16 bytes (32 hex chars)",
                    )
                channel_key_hex = request.channel_key.upper()
            except ValueError:
                return DecryptResult(
                    started=False,
                    total_packets=0,
                    message="Invalid hex string for channel key",
                )
        elif request.channel_name:
            # Derive key from channel name (hashtag channel)
            channel_key_bytes = sha256(request.channel_name.encode("utf-8")).digest()[:16]
            channel_key_hex = channel_key_bytes.hex().upper()
        else:
            return DecryptResult(
                started=False,
                total_packets=0,
                message="Must provide channel_key or channel_name",
            )
    else:
        # Contact decryption not yet supported (requires Ed25519 shared secret)
        return DecryptResult(
            started=False,
            total_packets=0,
            message="Contact key decryption not yet supported",
        )

    # Get count of undecrypted packets
    count = await RawPacketRepository.get_undecrypted_count()
    if count == 0:
        return DecryptResult(
            started=False, total_packets=0, message="No undecrypted packets to process"
        )

    # Start background decryption
    background_tasks.add_task(_run_historical_decryption, channel_key_bytes, channel_key_hex)

    return DecryptResult(
        started=True,
        total_packets=count,
        message=f"Started decryption of {count} packets in background",
    )


@router.get("/decrypt/progress", response_model=DecryptProgress | None)
async def get_decrypt_progress() -> DecryptProgress | None:
    """Get the current progress of historical decryption."""
    return _decrypt_progress