refactor(tasks): extract filename regex to shared constant
Move the valid filename regex pattern to a shared constant in `app/utils/filename_utils.py` and update both the task logic and security tests to use it. This eliminates duplication and ensures consistency across the codebase. Changes: - Defined `VALID_FILENAME_PATTERN` and `VALID_FILENAME_RE` in `app/utils/filename_utils.py`. - Updated `app/tasks/extract_metadata_with_gpt.py` to use `VALID_FILENAME_RE`. - Updated `tests/test_path_traversal_security.py` to use `VALID_FILENAME_PATTERN`. This refactoring addresses the duplication mentioned in the TODO in `tests/test_path_traversal_security.py`. Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
google-labs-jules[bot]
@@ -1,191 +1,192 @@
|
|||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
|
|
||||||
# Import the shared Celery instance
|
# Import the shared Celery instance
|
||||||
from app.celery_app import celery
|
from app.celery_app import celery
|
||||||
from app.config import settings
|
from app.config import settings
|
||||||
from app.database import SessionLocal
|
from app.database import SessionLocal
|
||||||
from app.models import FileRecord
|
from app.models import FileRecord
|
||||||
from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf
|
from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf
|
||||||
from app.tasks.retry_config import BaseTaskWithRetry
|
from app.tasks.retry_config import BaseTaskWithRetry
|
||||||
from app.utils import log_task_progress
|
from app.utils import log_task_progress
|
||||||
from app.utils.ai_provider import get_ai_provider
|
from app.utils.ai_provider import get_ai_provider
|
||||||
|
from app.utils.filename_utils import VALID_FILENAME_RE
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
def extract_json_from_text(text):
|
|
||||||
"""
|
def extract_json_from_text(text):
|
||||||
Try to extract a JSON object from the text.
|
"""
|
||||||
- First, check for a JSON block inside triple backticks.
|
Try to extract a JSON object from the text.
|
||||||
- If not found, try to extract text from the first '{' to the last '}'.
|
- First, check for a JSON block inside triple backticks.
|
||||||
"""
|
- If not found, try to extract text from the first '{' to the last '}'.
|
||||||
pattern = r"```(?:json)?\s*(\{.*?\})\s*```"
|
"""
|
||||||
match = re.search(pattern, text, re.DOTALL)
|
pattern = r"```(?:json)?\s*(\{.*?\})\s*```"
|
||||||
if match:
|
match = re.search(pattern, text, re.DOTALL)
|
||||||
return match.group(1)
|
if match:
|
||||||
else:
|
return match.group(1)
|
||||||
start = text.find("{")
|
else:
|
||||||
end = text.rfind("}")
|
start = text.find("{")
|
||||||
if start != -1 and end != -1 and end > start:
|
end = text.rfind("}")
|
||||||
return text[start : end + 1]
|
if start != -1 and end != -1 and end > start:
|
||||||
return None
|
return text[start : end + 1]
|
||||||
|
return None
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry, bind=True)
|
|
||||||
def extract_metadata_with_gpt(self, filename: str, cleaned_text: str, file_id: int = None):
|
@celery.task(base=BaseTaskWithRetry, bind=True)
|
||||||
"""
|
def extract_metadata_with_gpt(self, filename: str, cleaned_text: str, file_id: int = None):
|
||||||
Uses OpenAI to classify document metadata.
|
"""
|
||||||
|
Uses OpenAI to classify document metadata.
|
||||||
Args:
|
|
||||||
filename: Can be either a basename (e.g., "file.pdf") or a full path (e.g., "/workdir/processed/file.pdf")
|
Args:
|
||||||
cleaned_text: The extracted text from the document
|
filename: Can be either a basename (e.g., "file.pdf") or a full path (e.g., "/workdir/processed/file.pdf")
|
||||||
file_id: Optional file ID for tracking
|
cleaned_text: The extracted text from the document
|
||||||
"""
|
file_id: Optional file ID for tracking
|
||||||
task_id = self.request.id
|
"""
|
||||||
logger.info(f"[{task_id}] Starting metadata extraction for: {filename}")
|
task_id = self.request.id
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Starting metadata extraction for: {filename}")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"extract_metadata_with_gpt",
|
task_id,
|
||||||
"in_progress",
|
"extract_metadata_with_gpt",
|
||||||
f"Extracting metadata for {os.path.basename(filename)}",
|
"in_progress",
|
||||||
file_id=file_id,
|
f"Extracting metadata for {os.path.basename(filename)}",
|
||||||
)
|
file_id=file_id,
|
||||||
|
)
|
||||||
# Get file_id from database if not provided
|
|
||||||
if file_id is None:
|
# Get file_id from database if not provided
|
||||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
if file_id is None:
|
||||||
# Handle both basename and full path
|
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||||
if os.path.isabs(filename):
|
# Handle both basename and full path
|
||||||
file_path = filename
|
if os.path.isabs(filename):
|
||||||
else:
|
file_path = filename
|
||||||
file_path = os.path.join(tmp_dir, filename)
|
else:
|
||||||
if os.path.exists(file_path):
|
file_path = os.path.join(tmp_dir, filename)
|
||||||
with SessionLocal() as db:
|
if os.path.exists(file_path):
|
||||||
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
with SessionLocal() as db:
|
||||||
if file_record:
|
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
||||||
file_id = file_record.id
|
if file_record:
|
||||||
|
file_id = file_record.id
|
||||||
prompt = (
|
|
||||||
"You are a specialized document analyzer trained to extract structured metadata from documents.\n"
|
prompt = (
|
||||||
"Your task is to analyze the given text and return a well-structured JSON object.\n\n"
|
"You are a specialized document analyzer trained to extract structured metadata from documents.\n"
|
||||||
"Extract and return the following fields:\n"
|
"Your task is to analyze the given text and return a well-structured JSON object.\n\n"
|
||||||
"1. **filename**: Machine-readable filename "
|
"Extract and return the following fields:\n"
|
||||||
"(YYYY-MM-DD_DescriptiveTitle, use only letters, numbers, periods, and underscores).\n"
|
"1. **filename**: Machine-readable filename "
|
||||||
'2. **empfaenger**: The recipient, or "Unknown" if not found.\n'
|
"(YYYY-MM-DD_DescriptiveTitle, use only letters, numbers, periods, and underscores).\n"
|
||||||
'3. **absender**: The sender, or "Unknown" if not found.\n'
|
'2. **empfaenger**: The recipient, or "Unknown" if not found.\n'
|
||||||
"4. **correspondent**: The entity or company that issued the document "
|
'3. **absender**: The sender, or "Unknown" if not found.\n'
|
||||||
'(shortest possible name, e.g., "Amazon" instead of "Amazon EU SARL, German branch").\n'
|
"4. **correspondent**: The entity or company that issued the document "
|
||||||
"5. **kommunikationsart**: One of [Behoerdlicher_Brief, Rechnung, Kontoauszug, Vertrag, "
|
'(shortest possible name, e.g., "Amazon" instead of "Amazon EU SARL, German branch").\n'
|
||||||
"Quittung, Privater_Brief, Einladung, Gewerbliche_Korrespondenz, Newsletter, Werbung, Sonstiges].\n"
|
"5. **kommunikationsart**: One of [Behoerdlicher_Brief, Rechnung, Kontoauszug, Vertrag, "
|
||||||
"6. **kommunikationskategorie**: One of [Amtliche_Postbehoerdliche_Dokumente, "
|
"Quittung, Privater_Brief, Einladung, Gewerbliche_Korrespondenz, Newsletter, Werbung, Sonstiges].\n"
|
||||||
"Finanz_und_Vertragsdokumente, Geschaeftliche_Kommunikation, "
|
"6. **kommunikationskategorie**: One of [Amtliche_Postbehoerdliche_Dokumente, "
|
||||||
"Private_Korrespondenz, Sonstige_Informationen].\n"
|
"Finanz_und_Vertragsdokumente, Geschaeftliche_Kommunikation, "
|
||||||
"7. **document_type**: Precise classification (e.g., Invoice, Contract, Information, Unknown).\n"
|
"Private_Korrespondenz, Sonstige_Informationen].\n"
|
||||||
"8. **tags**: A list of up to 4 relevant thematic keywords.\n"
|
"7. **document_type**: Precise classification (e.g., Invoice, Contract, Information, Unknown).\n"
|
||||||
'9. **language**: Detected document language (ISO 639-1 code, e.g., "de" or "en").\n'
|
"8. **tags**: A list of up to 4 relevant thematic keywords.\n"
|
||||||
"10. **title**: A human-readable title summarizing the document content.\n"
|
'9. **language**: Detected document language (ISO 639-1 code, e.g., "de" or "en").\n'
|
||||||
"11. **confidence_score**: A numeric value (0-100) indicating the confidence level "
|
"10. **title**: A human-readable title summarizing the document content.\n"
|
||||||
"of the extracted metadata.\n"
|
"11. **confidence_score**: A numeric value (0-100) indicating the confidence level "
|
||||||
"12. **reference_number**: Extracted invoice/order/reference number if available.\n"
|
"of the extracted metadata.\n"
|
||||||
"13. **monetary_amounts**: A list of key monetary values detected in the document.\n\n"
|
"12. **reference_number**: Extracted invoice/order/reference number if available.\n"
|
||||||
"### Important Rules:\n"
|
"13. **monetary_amounts**: A list of key monetary values detected in the document.\n\n"
|
||||||
"- **OCR Correction**: Assume the text has been corrected for OCR errors.\n"
|
"### Important Rules:\n"
|
||||||
"- **Tagging**: Max 4 tags, avoiding generic or overly specific terms.\n"
|
"- **OCR Correction**: Assume the text has been corrected for OCR errors.\n"
|
||||||
"- **Title**: Concise, no addresses, and contains key identifying features.\n"
|
"- **Tagging**: Max 4 tags, avoiding generic or overly specific terms.\n"
|
||||||
"- **Date Selection**: Use the most relevant date if multiple are found.\n"
|
"- **Title**: Concise, no addresses, and contains key identifying features.\n"
|
||||||
"- **Output Language**: Maintain the document's original language.\n\n"
|
"- **Date Selection**: Use the most relevant date if multiple are found.\n"
|
||||||
f"Extracted text:\n{cleaned_text}\n\n"
|
"- **Output Language**: Maintain the document's original language.\n\n"
|
||||||
"Return only valid JSON with no additional commentary.\n"
|
f"Extracted text:\n{cleaned_text}\n\n"
|
||||||
)
|
"Return only valid JSON with no additional commentary.\n"
|
||||||
|
)
|
||||||
try:
|
|
||||||
logger.info(f"[{task_id}] Sending classification request for {filename}...")
|
try:
|
||||||
log_task_progress(task_id, "call_ai_provider", "in_progress", "Calling AI provider API", file_id=file_id)
|
logger.info(f"[{task_id}] Sending classification request for {filename}...")
|
||||||
provider = get_ai_provider()
|
log_task_progress(task_id, "call_ai_provider", "in_progress", "Calling AI provider API", file_id=file_id)
|
||||||
model = settings.ai_model or settings.openai_model
|
provider = get_ai_provider()
|
||||||
content = provider.chat_completion(
|
model = settings.ai_model or settings.openai_model
|
||||||
messages=[
|
content = provider.chat_completion(
|
||||||
{"role": "system", "content": "You are an intelligent document classifier."},
|
messages=[
|
||||||
{"role": "user", "content": prompt},
|
{"role": "system", "content": "You are an intelligent document classifier."},
|
||||||
],
|
{"role": "user", "content": prompt},
|
||||||
model=model,
|
],
|
||||||
temperature=0,
|
model=model,
|
||||||
)
|
temperature=0,
|
||||||
|
)
|
||||||
logger.info(f"[{task_id}] Raw classification response for {filename}: {content[:200]}...")
|
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Raw classification response for {filename}: {content[:200]}...")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"call_ai_provider",
|
task_id,
|
||||||
"success",
|
"call_ai_provider",
|
||||||
"Received AI provider response",
|
"success",
|
||||||
file_id=file_id,
|
"Received AI provider response",
|
||||||
detail=f"Raw classification response:\n{content}",
|
file_id=file_id,
|
||||||
)
|
detail=f"Raw classification response:\n{content}",
|
||||||
|
)
|
||||||
json_text = extract_json_from_text(content)
|
|
||||||
if not json_text:
|
json_text = extract_json_from_text(content)
|
||||||
logger.error(f"[{task_id}] Could not find valid JSON in GPT response for {filename}.")
|
if not json_text:
|
||||||
log_task_progress(
|
logger.error(f"[{task_id}] Could not find valid JSON in GPT response for {filename}.")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"extract_metadata_with_gpt",
|
task_id,
|
||||||
"failure",
|
"extract_metadata_with_gpt",
|
||||||
"Invalid JSON in response",
|
"failure",
|
||||||
file_id=file_id,
|
"Invalid JSON in response",
|
||||||
detail=f"Could not parse valid JSON from GPT response.\nRaw response:\n{content}",
|
file_id=file_id,
|
||||||
)
|
detail=f"Could not parse valid JSON from GPT response.\nRaw response:\n{content}",
|
||||||
return {}
|
)
|
||||||
|
return {}
|
||||||
metadata = json.loads(json_text)
|
|
||||||
|
metadata = json.loads(json_text)
|
||||||
# SECURITY: Validate filename format from GPT to prevent path traversal
|
|
||||||
# The prompt requests filenames with only letters, numbers, periods, and underscores
|
# SECURITY: Validate filename format from GPT to prevent path traversal
|
||||||
# Enforce this constraint to prevent malicious filenames
|
# The prompt requests filenames with only letters, numbers, periods, and underscores
|
||||||
suggested_filename = metadata.get("filename", "")
|
# Enforce this constraint to prevent malicious filenames
|
||||||
if suggested_filename:
|
suggested_filename = metadata.get("filename", "")
|
||||||
# Check if filename contains only safe characters AND explicitly check for ".."
|
if suggested_filename:
|
||||||
# Defense in depth: While the regex [\w\-\. ]+ already excludes / and \,
|
# Check if filename contains only safe characters AND explicitly check for ".."
|
||||||
# we explicitly reject ".." to guard against:
|
# Defense in depth: While the regex VALID_FILENAME_PATTERN already excludes / and \,
|
||||||
# 1. Potential locale-specific \w behavior
|
# we explicitly reject ".." to guard against:
|
||||||
# 2. Files literally named ".." which are valid but problematic
|
# 1. Potential locale-specific \w behavior
|
||||||
# 3. Future code changes that might relax the regex
|
# 2. Files literally named ".." which are valid but problematic
|
||||||
if not re.match(r"^[\w\-\. ]+$", suggested_filename) or ".." in suggested_filename:
|
# 3. Future code changes that might relax the regex
|
||||||
logger.warning(f"[{task_id}] Invalid filename format from GPT: '{suggested_filename}', using fallback")
|
if not VALID_FILENAME_RE.match(suggested_filename) or ".." in suggested_filename:
|
||||||
# Reset to empty to trigger fallback to original filename
|
logger.warning(f"[{task_id}] Invalid filename format from GPT: '{suggested_filename}', using fallback")
|
||||||
metadata["filename"] = ""
|
# Reset to empty to trigger fallback to original filename
|
||||||
|
metadata["filename"] = ""
|
||||||
logger.info(f"[{task_id}] Extracted metadata: {metadata}")
|
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Extracted metadata: {metadata}")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"parse_metadata",
|
task_id,
|
||||||
"success",
|
"parse_metadata",
|
||||||
f"Parsed metadata: {list(metadata.keys())}",
|
"success",
|
||||||
file_id=file_id,
|
f"Parsed metadata: {list(metadata.keys())}",
|
||||||
detail=f"Extracted metadata:\n{json.dumps(metadata, ensure_ascii=False, indent=2)}",
|
file_id=file_id,
|
||||||
)
|
detail=f"Extracted metadata:\n{json.dumps(metadata, ensure_ascii=False, indent=2)}",
|
||||||
|
)
|
||||||
# Trigger the next step: embedding metadata into the PDF
|
|
||||||
# Pass the filename (can be basename or full path) so embed_metadata_into_pdf can find the file on disk
|
# Trigger the next step: embedding metadata into the PDF
|
||||||
logger.info(f"[{task_id}] Queueing metadata embedding task")
|
# Pass the filename (can be basename or full path) so embed_metadata_into_pdf can find the file on disk
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Queueing metadata embedding task")
|
||||||
task_id, "extract_metadata_with_gpt", "success", "Metadata extracted, queuing embed task", file_id=file_id
|
log_task_progress(
|
||||||
)
|
task_id, "extract_metadata_with_gpt", "success", "Metadata extracted, queuing embed task", file_id=file_id
|
||||||
embed_metadata_into_pdf.delay(filename, cleaned_text, metadata, file_id)
|
)
|
||||||
|
embed_metadata_into_pdf.delay(filename, cleaned_text, metadata, file_id)
|
||||||
return {"s3_file": os.path.basename(filename), "metadata": metadata}
|
|
||||||
|
return {"s3_file": os.path.basename(filename), "metadata": metadata}
|
||||||
except Exception as e:
|
|
||||||
logger.exception(f"[{task_id}] AI provider classification failed for {filename}: {e}")
|
except Exception as e:
|
||||||
log_task_progress(
|
logger.exception(f"[{task_id}] AI provider classification failed for {filename}: {e}")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"extract_metadata_with_gpt",
|
task_id,
|
||||||
"failure",
|
"extract_metadata_with_gpt",
|
||||||
f"Exception: {str(e)}",
|
"failure",
|
||||||
file_id=file_id,
|
f"Exception: {str(e)}",
|
||||||
detail=f"AI provider classification failed for {filename}.\nException: {str(e)}",
|
file_id=file_id,
|
||||||
)
|
detail=f"AI provider classification failed for {filename}.\nException: {str(e)}",
|
||||||
return {}
|
)
|
||||||
|
return {}
|
||||||
|
|||||||
@@ -8,6 +8,11 @@ from pathlib import Path
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Pattern for valid filenames (alphanumeric, dash, underscore, period, and space)
|
||||||
|
# Used for validating GPT-provided filenames and other inputs
|
||||||
|
VALID_FILENAME_PATTERN = r"^[\w\-\. ]+$"
|
||||||
|
VALID_FILENAME_RE = re.compile(VALID_FILENAME_PATTERN)
|
||||||
|
|
||||||
|
|
||||||
def get_unique_filename(original_path: str, check_exists_func: Callable[[str], bool] | None = None) -> str:
|
def get_unique_filename(original_path: str, check_exists_func: Callable[[str], bool] | None = None) -> str:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -212,9 +212,10 @@ class TestExtractMetadataFilenameValidation:
|
|||||||
"""Test that invalid filename formats are rejected."""
|
"""Test that invalid filename formats are rejected."""
|
||||||
import re
|
import re
|
||||||
|
|
||||||
# Valid pattern from extract_metadata_with_gpt.py
|
from app.utils.filename_utils import VALID_FILENAME_PATTERN
|
||||||
# TODO: Consider extracting this to a shared constant to avoid duplication
|
|
||||||
valid_pattern = r"^[\w\-\. ]+$"
|
# Valid pattern from app.utils.filename_utils
|
||||||
|
valid_pattern = VALID_FILENAME_PATTERN
|
||||||
|
|
||||||
# Test valid filenames
|
# Test valid filenames
|
||||||
valid_filenames = [
|
valid_filenames = [
|
||||||
|
|||||||
Reference in New Issue
Block a user