mirror of
https://github.com/dobin/detonator
synced 2026-08-09 12:31:13 +00:00
refactor: store files on fs not in db
This commit is contained in:
@@ -8,3 +8,4 @@ config.env
|
||||
.vscode/
|
||||
*.yaml
|
||||
todo.md
|
||||
upload/
|
||||
@@ -4,9 +4,11 @@ from typing import Optional, Dict
|
||||
import time
|
||||
import json
|
||||
from datetime import datetime
|
||||
from detonatorapi.edr_parser.edr_parser import EdrParser
|
||||
from typing import List
|
||||
import os
|
||||
|
||||
from detonatorapi.edr_parser.edr_parser import EdrParser
|
||||
from detonatorapi.settings import UPLOAD_DIR
|
||||
from detonatorapi.database import Scan, get_db_for_thread
|
||||
from detonatorapi.db_interface import db_scan_change_status_quick, db_scan_add_log
|
||||
from detonatorapi.agent.agent_api import AgentApi
|
||||
@@ -82,7 +84,12 @@ def scan_file_with_agent(scan_id: int) -> bool:
|
||||
|
||||
filename = db_scan.file.filename
|
||||
exec_arguments = db_scan.file.exec_arguments
|
||||
file_content = db_scan.file.content
|
||||
|
||||
# Read file content from disk
|
||||
file_path = os.path.join(UPLOAD_DIR, db_scan.file.filename)
|
||||
with open(file_path, 'rb') as f:
|
||||
file_content = f.read()
|
||||
|
||||
runtime = db_scan.runtime
|
||||
drop_path = db_scan.drop_path
|
||||
rededr_port = db_scan.profile.rededr_port
|
||||
|
||||
@@ -38,7 +38,6 @@ class File(Base):
|
||||
__tablename__ = "files"
|
||||
|
||||
id: Mapped[int] = Column(Integer, primary_key=True, index=True)
|
||||
content: Mapped[bytes] = Column(LargeBinary, nullable=False)
|
||||
filename: Mapped[str] = Column(String(255), nullable=False)
|
||||
exec_arguments: Mapped[str] = Column(String(255), nullable=True)
|
||||
file_hash: Mapped[str] = Column(String(64), nullable=False, index=True)
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
from typing import Optional, List
|
||||
from datetime import datetime
|
||||
import logging
|
||||
import os
|
||||
from .settings import UPLOAD_DIR
|
||||
import random
|
||||
import string
|
||||
from werkzeug.utils import secure_filename
|
||||
|
||||
from .database import Scan, File, Profile, get_db_for_thread
|
||||
from .utils import mylog
|
||||
@@ -48,11 +53,20 @@ def db_scan_add_log(db, db_scan, log_message: str):
|
||||
|
||||
def db_create_file(db, filename: str, content: bytes, source_url: str = "", comment: str = "", exec_arguments: str = "", user: str = "") -> int:
|
||||
file_hash = File.calculate_hash(content)
|
||||
|
||||
# prepend 4 random chars to filename to avoid collisions
|
||||
filename = secure_filename(filename)
|
||||
rand_str = ''.join(random.choices(string.ascii_letters + string.digits, k=4))
|
||||
actual_filename = f"{rand_str}_{filename}"
|
||||
|
||||
# DB: Create file record
|
||||
# Write file content to disk
|
||||
file_path = os.path.join(UPLOAD_DIR, f"{actual_filename}")
|
||||
with open(file_path, 'wb') as f:
|
||||
f.write(content)
|
||||
|
||||
# DB: Create file record with path instead of content
|
||||
db_file = File(
|
||||
content=content,
|
||||
filename=filename,
|
||||
filename=actual_filename,
|
||||
file_hash=file_hash,
|
||||
source_url=source_url,
|
||||
comment=comment,
|
||||
@@ -62,7 +76,7 @@ def db_create_file(db, filename: str, content: bytes, source_url: str = "", comm
|
||||
db.add(db_file)
|
||||
db.commit()
|
||||
|
||||
logger.info(f"DB: Created file {db_file.id} with filename: {filename}")
|
||||
logger.info(f"DB: Created file {db_file.id} with filename: {actual_filename}")
|
||||
return db_file.id
|
||||
|
||||
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import string
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File as FastAPIFile, Form
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
@@ -134,12 +132,9 @@ async def upload_file_and_scan(
|
||||
filename = file.filename
|
||||
if not filename:
|
||||
raise HTTPException(status_code=400, detail="Filename cannot be empty")
|
||||
rand_str = ''.join(random.choices(string.ascii_letters + string.digits, k=4))
|
||||
actual_filename = f"{rand_str}_{filename}"
|
||||
logger.info(f"Uploading file: {actual_filename}")
|
||||
file_content = await file.read()
|
||||
file_id = db_create_file(db,
|
||||
filename=actual_filename,
|
||||
filename=filename,
|
||||
content=file_content,
|
||||
source_url=source_url or "",
|
||||
comment=file_comment or "",
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
|
||||
UPLOAD_DIR = "upload/"
|
||||
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
|
||||
@@ -3,11 +3,13 @@ from fastapi.responses import Response
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
from typing import List, Optional
|
||||
import logging
|
||||
import os
|
||||
|
||||
from .database import get_db, File, Scan
|
||||
from .schemas import FileResponse, FileWithScans
|
||||
from .db_interface import db_create_file, db_create_scan, db_get_profile_by_name
|
||||
from .token_auth import require_auth, get_user_from_request
|
||||
from .settings import UPLOAD_DIR
|
||||
|
||||
router = APIRouter()
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -75,6 +77,15 @@ async def delete_file(
|
||||
if db_file is None:
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
|
||||
# Delete file from filesystem
|
||||
file_path = os.path.join(UPLOAD_DIR, db_file.filename)
|
||||
if os.path.exists(file_path):
|
||||
try:
|
||||
os.remove(file_path)
|
||||
logger.info(f"Deleted file from disk: {file_path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete file from disk: {file_path}, error: {e}")
|
||||
|
||||
# Delete associated scans first
|
||||
db.query(Scan).filter(Scan.file_id == file_id).delete()
|
||||
db.delete(db_file)
|
||||
@@ -93,8 +104,17 @@ async def download_file(
|
||||
if db_file is None:
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
|
||||
# Check if file exists on disk
|
||||
file_path = os.path.join(UPLOAD_DIR, db_file.filename)
|
||||
if not os.path.exists(file_path):
|
||||
raise HTTPException(status_code=404, detail="File not found on disk")
|
||||
|
||||
# Read file content from disk
|
||||
with open(file_path, 'rb') as f:
|
||||
content = f.read()
|
||||
|
||||
return Response(
|
||||
content=db_file.content,
|
||||
content=content,
|
||||
media_type="application/octet-stream",
|
||||
headers={
|
||||
"Content-Disposition": f"attachment; filename={db_file.filename}"
|
||||
|
||||
Reference in New Issue
Block a user