dtf-system/app/scanning.py

"""Releasing artwork: ClamAV up to its size limit, a format check above it.

Unknown or error results NEVER release artwork. ClamAV scans files up to
scan_limit_bytes() (2 GB). Sheets of several GB are the normal order and
ClamAV cannot take them, so a larger file is released only if its first bytes
are those of the format its name claims (a PNG that really is a PNG, not a
program renamed .png). That is the check the client chose for large files; it
does not look for malware inside a valid file.
"""
import re
import socket
import struct
import time
from fastapi import HTTPException
from .core.auth import audit
from .core.db import connect
from .core.limits import scan_limit_bytes

def require_clean(row):
    if not row['complete'] or row['scan_state'] != 'clean':
        raise HTTPException(409, 'Artwork is quarantined until the malware scan succeeds')

class ClamAV:
    def command(self, command, timeout=2):
        with socket.create_connection(('scanner',3310),timeout=timeout) as sock:
            sock.settimeout(timeout)
            sock.sendall(b'z'+command+b'\0')
            reply=b''
            while b'\0' not in reply and len(reply)<4096:
                block=sock.recv(4096)
                if not block:break
                reply+=block
        return reply.rstrip(b'\0\n')

    def ping(self):
        return self.command(b'PING') == b'PONG'

    def version(self):
        return self.command(b'VERSION').decode('utf-8','replace')

    def scan(self, stream, size):
        with socket.create_connection(('scanner',3310),timeout=10) as sock:
            # A 2 GB file takes minutes to stream and scan.
            sock.settimeout(900)
            sock.sendall(b'zINSTREAM\0')
            sent=0
            for chunk in stream.iter_chunks(chunk_size=65536):
                sent+=len(chunk)
                if sent>size: return 'rejected','Stored file size changed'
                sock.sendall(struct.pack('!I',len(chunk))+chunk)
            if sent!=size:return 'rejected','Stored file size changed'
            sock.sendall(b'\0\0\0\0')
            reply=b''
            while b'\0' not in reply and len(reply)<4096:
                block=sock.recv(4096)
                if not block:break
                reply+=block
            result=reply.rstrip(b'\0\n')
            if result==b'stream: OK':return 'clean',None
            if result.endswith(b' FOUND'):return 'rejected','Malware or unsafe scan condition detected'
            return 'error','Scanner could not verify this file'

# What each accepted extension must start with. AI files are PDF or PostScript;
# CDR and WebP are RIFF containers with their own form type.
TIFF = (b'II*\x00', b'MM\x00*', b'II+\x00', b'MM\x00+')
SIGNATURES = {
    'png': (b'\x89PNG\r\n\x1a\n',),
    'jpg': (b'\xff\xd8\xff',), 'jpeg': (b'\xff\xd8\xff',),
    'tif': TIFF, 'tiff': TIFF,
    'pdf': (b'%PDF-',), 'ai': (b'%PDF-', b'%!PS'),
    'psd': (b'8BPS',), 'psb': (b'8BPS',),
}
RIFF_FORMS = {'cdr': re.compile(rb'^RIFF....CDR', re.S), 'webp': re.compile(rb'^RIFF....WEBP', re.S)}
HEAD_BYTES = 64


def format_matches(name, head):
    """Whether a file's first bytes are those of the format its name claims."""
    ext = name.rsplit('.', 1)[-1].lower() if '.' in name else ''
    if ext in RIFF_FORMS:
        return bool(RIFF_FORMS[ext].match(head))
    return any(head.startswith(sig) for sig in SIGNATURES.get(ext, ()))


def check_large(storage, row):
    """Release decision for a file above the antivirus limit."""
    head = storage.client.get_object(Bucket=storage.bucket, Key=row['object_key'],
                                     Range=f'bytes=0-{HEAD_BYTES - 1}')['Body'].read()
    if format_matches(row['name'], head):
        return 'clean', 'Acima do limite do antivírus; formato do arquivo conferido'
    return 'rejected', 'O conteúdo do arquivo não corresponde ao formato do nome'


def scan_one(storage, scanner=None):
    scanner=scanner or ClamAV()
    with connect() as c:
        row=c.execute('''SELECT * FROM dtf_local.uploads WHERE complete AND purged_at IS NULL
            AND expires_at>now() AND scan_state IN ('pending','error') AND scan_after<=now()
            ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1''').fetchone()
        if not row:return False
        try:
            if row['size'] > scan_limit_bytes():
                state,reason=check_large(storage,row)
            else:
                stream=storage.client.get_object(Bucket=storage.bucket,Key=row['object_key'])['Body']
                try:state,reason=scanner.scan(stream,row['size'])
                finally:stream.close()
        except Exception:
            state,reason='error','Malware scanner unavailable; file remains blocked'
        c.execute("""UPDATE dtf_local.uploads SET scan_state=%s,scan_reason=%s,scanned_at=now(),
            scan_after=now()+interval '1 minute',
            expires_at=CASE WHEN %s IN ('rejected','error') THEN LEAST(expires_at,now()+interval '3 days') ELSE expires_at END
            WHERE id=%s""",(state,reason,state,row['id']))
        audit('artwork_scan', upload=str(row['id']), result=state)
    return True

def scan_loop(storage):
    while True:
        try:
            if scan_one(storage):continue
        except Exception:
            audit('scanner_worker_error')
        time.sleep(1)