first commit

local/scanning.py

@@ -0,0 +1,81 @@
+"""Local ClamAV boundary. Unknown/error/over-limit results NEVER release artwork."""
+import os
+import socket
+import struct
+import time
+from fastapi import HTTPException
+from .auth import audit
+from .db import connect
+
+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):
+        if size > min(134217728, int(os.environ.get('SCAN_MAX_BYTES','134217728'))):
+            return 'rejected', 'File exceeds the local malware scan limit'
+        with socket.create_connection(('scanner',3310),timeout=10) as sock:
+            sock.settimeout(150)
+            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'
+
+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:
+            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','Local 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)