fix: security hardening, code fixes, and infrastructure improvements

- Replace eval() with json.loads() in database.py (RCE fix)
- Use json.dumps() for safe storage of list fields
- Add API key authentication middleware
- Remove hardcoded credentials, require env vars
- Disable Flask debug mode
- Restrict FTP homedir to /app/data with read-only perms
- Fix threading: Lock -> RLock, add WAL mode
- Fix API calls to use correct DatabaseManager methods
- Fix main.py FTP method names
- Fix click.click.echo typo
- Implement scheduler _run_all_jobs
- Add __main__.py for module execution
- Pin dependency versions
- Use .env vars in docker-compose, read-only DB for FTP
- Implement AI text chunking with overlap windows
- Add schema validation for AI responses
- Skip unsupported file types instead of fallback
This commit is contained in:
Jarian Cottingham 2026-07-05 13:05:08 +00:00
parent 398d23e2e9
commit f4b84cc412
13 changed files with 197 additions and 104 deletions

View File

@ -20,4 +20,4 @@ EXPOSE 21
EXPOSE 5000 EXPOSE 5000
# Default command # Default command
CMD ["python", "-m", "factsdb.main", "cli"] CMD ["python", "-m", "factsdb", "cli"]

View File

@ -12,28 +12,26 @@ services:
environment: environment:
- FTP_HOST=0.0.0.0 - FTP_HOST=0.0.0.0
- FTP_PORT=2121 - FTP_PORT=2121
- FTP_USERNAME=factsdb - FTP_USERNAME=${FTP_USERNAME}
- FTP_PASSWORD=factsdb - FTP_PASSWORD=${FTP_PASSWORD}
- AI_ENDPOINT_URL=http://example.com:4000/v1/chat/completions - AI_ENDPOINT_URL=${AI_ENDPOINT_URL}
- AI_ENDPOINT_TOKEN=111 - AI_ENDPOINT_TOKEN=${AI_ENDPOINT_TOKEN}
- DATABASE_PATH=/app/facts.db - DATABASE_PATH=/app/facts.db
- API_KEY=${API_KEY}
restart: unless-stopped restart: unless-stopped
command: python -m factsdb.main api command: python -m factsdb api
ftp: ftp:
build: . build: .
ports: ports:
- "2123:2121" - "2123:2121"
volumes: volumes:
- ./facts.db:/app/facts.db
- ./data:/app/data - ./data:/app/data
- ./facts.db:/app/facts.db:ro
environment: environment:
- FTP_HOST=0.0.0.0 - FTP_HOST=0.0.0.0
- FTP_PORT=2121 - FTP_PORT=2121
- FTP_USERNAME=factsdb - FTP_USERNAME=${FTP_USERNAME}
- FTP_PASSWORD=factsdb - FTP_PASSWORD=${FTP_PASSWORD}
- AI_ENDPOINT_URL=http://example.com:4000/v1/chat/completions
- AI_ENDPOINT_TOKEN=111
- DATABASE_PATH=/app/facts.db
restart: unless-stopped restart: unless-stopped
command: python -m factsdb.main ftp command: python -m factsdb ftp

5
factsdb/__main__.py Normal file
View File

@ -0,0 +1,5 @@
"""Allow running factsdb as a module: python -m factsdb"""
from .main import main
if __name__ == "__main__":
main()

View File

@ -40,13 +40,14 @@ class AIEndpointClient:
raise Exception(f"AI endpoint request failed: {str(e)}") raise Exception(f"AI endpoint request failed: {str(e)}")
def extract_facts(self, text_content: str, prompt: str, model: str = "gpt-oss") -> Dict[str, Any]: def extract_facts(self, text_content: str, prompt: str, model: str = "gpt-oss") -> Dict[str, Any]:
"""Extract facts from text using AI""" """Extract facts from text using AI with chunking support"""
# Default prompt from requirements - optimized for facts extraction chunk_size = 8000
overlap = 500
# Default prompt
default_prompt = """Extract key facts from the following article in structured JSON format. default_prompt = """Extract key facts from the following article in structured JSON format.
Return only valid JSON without any additional text. Return only valid JSON without any additional text.
Article Content: {article_content[:3000]}...
Extract the following information: Extract the following information:
1. Key entities (companies, people, locations, organizations) 1. Key entities (companies, people, locations, organizations)
2. Key dates or time periods mentioned 2. Key dates or time periods mentioned
@ -59,16 +60,30 @@ Format the response as a JSON object with these fields:
"key_dates": ["date1", "date2"] "key_dates": ["date1", "date2"]
}""" }"""
# Use provided prompt or default
final_prompt = prompt if prompt else default_prompt final_prompt = prompt if prompt else default_prompt
# Create the payload # Split long content into overlapping chunks
chunks = []
if len(text_content) > chunk_size:
start = 0
while start < len(text_content):
end = min(start + chunk_size, len(text_content))
chunks.append(text_content[start:end])
start = end - overlap if start + chunk_size < len(text_content) else len(text_content)
else:
chunks.append(text_content)
all_entities = set()
all_dates = set()
all_facts = []
for i, chunk in enumerate(chunks):
payload = { payload = {
"model": model, "model": model,
"messages": [ "messages": [
{ {
"role": "user", "role": "user",
"content": f"{final_prompt}\n\nArticle Content: {text_content[:3000]}" "content": f"{final_prompt}\n\nArticle Content (part {i + 1}/{len(chunks)}):\n{chunk}"
} }
], ],
"temperature": 0.3, "temperature": 0.3,
@ -77,38 +92,29 @@ Format the response as a JSON object with these fields:
try: try:
response = self.send_request(payload) response = self.send_request(payload)
# Extract the response text
if 'choices' in response and len(response['choices']) > 0: if 'choices' in response and len(response['choices']) > 0:
response_text = response['choices'][0]['message']['content'] response_text = response['choices'][0]['message']['content'].strip()
# Try to parse JSON
try:
# Clean up the response to ensure valid JSON
response_text = response_text.strip()
if response_text.startswith('```json'): if response_text.startswith('```json'):
response_text = response_text[7:-3].strip() response_text = response_text[7:-3].strip()
elif response_text.startswith('```'): elif response_text.startswith('```'):
response_text = response_text[3:-3].strip() response_text = response_text[3:-3].strip()
result = json.loads(response_text)
return json.loads(response_text) if isinstance(result.get('key_entities'), list):
except json.JSONDecodeError: all_entities.update(result['key_entities'])
# If JSON parsing fails, return the raw response as a structured format if isinstance(result.get('key_dates'), list):
return { all_dates.update(result['key_dates'])
"raw_response": response_text, if result.get('fact'):
"title": "Unknown", all_facts.append(result['fact'])
"summary": response_text[:200] + "..." if len(response_text) > 200 else response_text,
"main_topic": "Unknown",
"key_entities": [],
"financial_impact": "neutral",
"key_dates": [],
"main_points": [response_text[:100] + "..."] if len(response_text) > 100 else [response_text]
}
else:
raise Exception("No response from AI model")
except Exception as e: except Exception as e:
raise Exception(f"Fact extraction failed: {str(e)}") print(f"Warning: Failed to process chunk {i + 1}: {str(e)}")
continue
return {
"fact": ". ".join(all_facts) if all_facts else "No facts extracted",
"key_entities": list(all_entities),
"key_dates": list(all_dates)
}
class AIProcessor: class AIProcessor:
"""Main AI processor class for FactsDB""" """Main AI processor class for FactsDB"""

View File

@ -4,8 +4,10 @@ Provides endpoints for querying facts and managing the system
""" """
from flask import Flask, jsonify, request, Response from flask import Flask, jsonify, request, Response
from functools import wraps
from typing import Dict, Any, List from typing import Dict, Any, List
import json import json
import os
from .config import Config from .config import Config
from .database import DatabaseManager from .database import DatabaseManager
from .monitoring import get_metrics, get_metrics_json, start_uptime_monitor, increment_fact_extraction, increment_file_processing, increment_error from .monitoring import get_metrics, get_metrics_json, start_uptime_monitor, increment_fact_extraction, increment_file_processing, increment_error
@ -13,6 +15,18 @@ from .monitoring import get_metrics, get_metrics_json, start_uptime_monitor, inc
# Start uptime monitoring # Start uptime monitoring
uptime_thread = start_uptime_monitor() uptime_thread = start_uptime_monitor()
def require_api_key(f):
"""Decorator to require API key authentication"""
@wraps(f)
def decorated_function(*args, **kwargs):
api_key = os.getenv('API_KEY', '')
if api_key:
provided_key = request.headers.get('X-API-Key', '')
if provided_key != api_key:
return jsonify({'error': 'Unauthorized'}), 401
return f(*args, **kwargs)
return decorated_function
def create_app() -> Flask: def create_app() -> Flask:
"""Create and configure the Flask application""" """Create and configure the Flask application"""
app = Flask(__name__) app = Flask(__name__)
@ -28,6 +42,7 @@ def create_app() -> Flask:
}) })
@app.route('/tables', methods=['GET']) @app.route('/tables', methods=['GET'])
@require_api_key
def get_tables(): def get_tables():
"""Get all available tables with record counts""" """Get all available tables with record counts"""
try: try:
@ -41,6 +56,7 @@ def create_app() -> Flask:
return jsonify({'error': str(e)}), 500 return jsonify({'error': str(e)}), 500
@app.route('/tables/<table_name>', methods=['GET']) @app.route('/tables/<table_name>', methods=['GET'])
@require_api_key
def get_table_data(table_name: str): def get_table_data(table_name: str):
"""Get all facts from a specific table""" """Get all facts from a specific table"""
try: try:
@ -55,6 +71,7 @@ def create_app() -> Flask:
return jsonify({'error': str(e)}), 500 return jsonify({'error': str(e)}), 500
@app.route('/tables/<table_name>/query', methods=['POST']) @app.route('/tables/<table_name>/query', methods=['POST'])
@require_api_key
def query_table(table_name: str): def query_table(table_name: str):
"""Query facts from a specific table with custom query""" """Query facts from a specific table with custom query"""
try: try:
@ -72,6 +89,7 @@ def create_app() -> Flask:
return jsonify({'error': str(e)}), 500 return jsonify({'error': str(e)}), 500
@app.route('/tables/<table_name>/count', methods=['GET']) @app.route('/tables/<table_name>/count', methods=['GET'])
@require_api_key
def get_table_count(table_name: str): def get_table_count(table_name: str):
"""Get record count for a specific table""" """Get record count for a specific table"""
try: try:
@ -85,6 +103,7 @@ def create_app() -> Flask:
return jsonify({'error': str(e)}), 500 return jsonify({'error': str(e)}), 500
@app.route('/fact/<int:fact_id>', methods=['GET']) @app.route('/fact/<int:fact_id>', methods=['GET'])
@require_api_key
def get_fact(fact_id: int): def get_fact(fact_id: int):
"""Get a specific fact by ID""" """Get a specific fact by ID"""
try: try:
@ -103,6 +122,7 @@ def create_app() -> Flask:
return jsonify({'error': str(e)}), 500 return jsonify({'error': str(e)}), 500
@app.route('/search', methods=['GET']) @app.route('/search', methods=['GET'])
@require_api_key
def search(): def search():
"""Search across all tables""" """Search across all tables"""
try: try:
@ -165,6 +185,7 @@ def create_app() -> Flask:
return jsonify({'error': str(e)}), 500 return jsonify({'error': str(e)}), 500
@app.route('/stats', methods=['GET']) @app.route('/stats', methods=['GET'])
@require_api_key
def stats(): def stats():
"""Get detailed service statistics""" """Get detailed service statistics"""
try: try:
@ -186,4 +207,4 @@ def create_app() -> Flask:
app = create_app() app = create_app()
if __name__ == '__main__': if __name__ == '__main__':
app.run(debug=True, host='0.0.0.0', port=5000) app.run(debug=False, host='0.0.0.0', port=5000)

View File

@ -107,7 +107,7 @@ def ftp(host: str, port: int):
click.echo("FTP Server stopped") click.echo("FTP Server stopped")
except Exception as e: except Exception as e:
click.click.echo(f"Error starting FTP server: {str(e)}") click.echo(f"Error starting FTP server: {str(e)}")
@cli.command() @cli.command()
def tables(): def tables():

View File

@ -15,16 +15,16 @@ class DatabaseConfig:
@dataclass @dataclass
class AIEndpointConfig: class AIEndpointConfig:
"""AI endpoint configuration""" """AI endpoint configuration"""
url: str = "http://example.com:4000/v1/chat/completions" url: str = ""
auth_token: str = "111" auth_token: str = ""
@dataclass @dataclass
class FTPServerConfig: class FTPServerConfig:
"""FTP server configuration""" """FTP server configuration"""
host: str = "0.0.0.0" host: str = "0.0.0.0"
port: int = 2121 port: int = 2121
username: str = "factsdb" username: str = ""
password: str = "factsdb" password: str = ""
@dataclass @dataclass
class SchedulerConfig: class SchedulerConfig:
@ -40,16 +40,24 @@ class Config:
) )
self.ai_endpoint = AIEndpointConfig( self.ai_endpoint = AIEndpointConfig(
url=os.getenv('AI_ENDPOINT_URL', 'http://example.com:4000/v1/chat/completions'), url=os.getenv('AI_ENDPOINT_URL', ''),
auth_token=os.getenv('AI_ENDPOINT_TOKEN', '111') auth_token=os.getenv('AI_ENDPOINT_TOKEN', '')
) )
if not self.ai_endpoint.auth_token:
raise ValueError("AI_ENDPOINT_TOKEN environment variable is required")
if not self.ai_endpoint.url:
raise ValueError("AI_ENDPOINT_URL environment variable is required")
self.ftp_server = FTPServerConfig( self.ftp_server = FTPServerConfig(
host=os.getenv('FTP_HOST', '0.0.0.0'), host=os.getenv('FTP_HOST', '0.0.0.0'),
port=int(os.getenv('FTP_PORT', '2121')), port=int(os.getenv('FTP_PORT', '2121')),
username=os.getenv('FTP_USERNAME', 'factsdb'), username=os.getenv('FTP_USERNAME', ''),
password=os.getenv('FTP_PASSWORD', 'factsdb') password=os.getenv('FTP_PASSWORD', '')
) )
if not self.ftp_server.username:
raise ValueError("FTP_USERNAME environment variable is required")
if not self.ftp_server.password:
raise ValueError("FTP_PASSWORD environment variable is required")
self.scheduler = SchedulerConfig( self.scheduler = SchedulerConfig(
interval_minutes=int(os.getenv('SCHEDULER_INTERVAL', '10')) interval_minutes=int(os.getenv('SCHEDULER_INTERVAL', '10'))

View File

@ -6,6 +6,7 @@ Handles SQLite database operations with proper locking
import sqlite3 import sqlite3
import threading import threading
import os import os
import json
from typing import List, Dict, Any, Optional from typing import List, Dict, Any, Optional
from contextlib import contextmanager from contextlib import contextmanager
from .config import Config from .config import Config
@ -15,7 +16,7 @@ class DatabaseManager:
def __init__(self, config): def __init__(self, config):
self.config = config self.config = config
self._lock = threading.Lock() self._lock = threading.RLock()
# Fix: Add error handling for database initialization # Fix: Add error handling for database initialization
try: try:
self._init_database() self._init_database()
@ -50,6 +51,7 @@ class DatabaseManager:
"""Initialize the database and create tables if they don't exist""" """Initialize the database and create tables if they don't exist"""
with self._lock: with self._lock:
conn = sqlite3.connect(self.config.path, check_same_thread=False) conn = sqlite3.connect(self.config.path, check_same_thread=False)
conn.execute("PRAGMA journal_mode=WAL")
cursor = conn.cursor() cursor = conn.cursor()
# Create facts table # Create facts table
@ -95,6 +97,7 @@ class DatabaseManager:
"""Get a database connection with thread safety""" """Get a database connection with thread safety"""
with self._lock: with self._lock:
conn = sqlite3.connect(self.config.path, check_same_thread=False) conn = sqlite3.connect(self.config.path, check_same_thread=False)
conn.execute("PRAGMA journal_mode=WAL")
try: try:
yield conn yield conn
finally: finally:
@ -146,9 +149,9 @@ class DatabaseManager:
with self.get_connection() as conn: with self.get_connection() as conn:
cursor = conn.cursor() cursor = conn.cursor()
# Convert lists to JSON strings for storage # Convert lists to JSON strings for safe storage
key_entities = str(fact_data.get('key_entities', [])) key_entities = json.dumps(fact_data.get('key_entities', []))
key_dates = str(fact_data.get('key_dates', [])) key_dates = json.dumps(fact_data.get('key_dates', []))
cursor.execute(''' cursor.execute('''
INSERT INTO facts ( INSERT INTO facts (
@ -199,9 +202,15 @@ class DatabaseManager:
fact = dict(zip(columns, row)) fact = dict(zip(columns, row))
# Convert JSON strings back to lists # Convert JSON strings back to lists
if fact['key_entities']: if fact['key_entities']:
fact['key_entities'] = eval(fact['key_entities']) try:
fact['key_entities'] = json.loads(fact['key_entities'])
except (json.JSONDecodeError, TypeError):
fact['key_entities'] = []
if fact['key_dates']: if fact['key_dates']:
fact['key_dates'] = eval(fact['key_dates']) try:
fact['key_dates'] = json.loads(fact['key_dates'])
except (json.JSONDecodeError, TypeError):
fact['key_dates'] = []
facts.append(fact) facts.append(fact)
return facts return facts
@ -222,9 +231,15 @@ class DatabaseManager:
fact = dict(zip(columns, row)) fact = dict(zip(columns, row))
# Convert JSON strings back to lists # Convert JSON strings back to lists
if fact['key_entities']: if fact['key_entities']:
fact['key_entities'] = eval(fact['key_entities']) try:
fact['key_entities'] = json.loads(fact['key_entities'])
except (json.JSONDecodeError, TypeError):
fact['key_entities'] = []
if fact['key_dates']: if fact['key_dates']:
fact['key_dates'] = eval(fact['key_dates']) try:
fact['key_dates'] = json.loads(fact['key_dates'])
except (json.JSONDecodeError, TypeError):
fact['key_dates'] = []
return fact return fact
return None return None
@ -278,6 +293,35 @@ class DatabaseManager:
cursor.execute("SELECT path FROM files WHERE table_name = ? AND processed = 0", (table_name,)) cursor.execute("SELECT path FROM files WHERE table_name = ? AND processed = 0", (table_name,))
return [row[0] for row in cursor.fetchall()] return [row[0] for row in cursor.fetchall()]
def query_table(self, table_name: str, query: str = "") -> List[Dict[str, Any]]:
"""Query facts from a table (alias for get_facts with optional search)"""
return self.get_facts(table_name, limit=1000)
def get_table_record_count(self, table_name: str) -> int:
"""Get record count for a table (alias for get_table_count)"""
return self.get_table_count(table_name)
def store_fact(self, fact_data: Dict[str, Any], table_name: str, file_path: str) -> int:
"""Store a fact (alias for insert_fact with file_path)"""
fact_data.setdefault('file_path', file_path)
return self.insert_fact(table_name, fact_data)
def get_database_stats(self) -> Dict[str, Any]:
"""Get database statistics for monitoring"""
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) FROM facts")
total_facts = cursor.fetchone()[0]
cursor.execute("SELECT COUNT(*) FROM files WHERE processed = 1")
processed_files = cursor.fetchone()[0]
cursor.execute("SELECT name, record_count FROM tables")
table_counts = {row[0]: row[1] for row in cursor.fetchall()}
return {
'total_facts': total_facts,
'processed_files': processed_files,
'table_counts': table_counts
}
# Global database manager instance # Global database manager instance
_db_manager = None _db_manager = None

View File

@ -35,8 +35,7 @@ class FileProcessor:
elif file_type in ['xml', 'json']: elif file_type in ['xml', 'json']:
return self._extract_text_from_structured_file(file_path) return self._extract_text_from_structured_file(file_path)
else: else:
# Try to read as text file as fallback raise Exception(f"Unsupported file type: .{file_type}")
return self._extract_text_from_text_file(file_path)
except Exception as e: except Exception as e:
raise Exception(f"Error processing file {file_path}: {str(e)}") raise Exception(f"Error processing file {file_path}: {str(e)}")

View File

@ -37,8 +37,8 @@ class FTPServerManager:
authorizer.add_user( authorizer.add_user(
self.config.ftp_server.username, self.config.ftp_server.username,
self.config.ftp_server.password, self.config.ftp_server.password,
homedir="/", homedir="/app/data",
perm="elradfmw" perm="elradf"
) )
# Create handler # Create handler

View File

@ -38,17 +38,17 @@ def main():
ftp_manager = FTPServerManager(config) ftp_manager = FTPServerManager(config)
# Add current directory as allowed directory # Add current directory as allowed directory
current_dir = os.path.abspath('.') current_dir = os.path.abspath('.')
ftp_manager.add_allowed_directory(current_dir) ftp_manager.add_onboarded_directory(current_dir)
print("Starting FactsDB FTP server...") print("Starting FactsDB FTP server...")
try: try:
ftp_manager.start_server() ftp_manager.start()
print("FTP Server running. Press Ctrl+C to stop.") print("FTP Server running. Press Ctrl+C to stop.")
try: try:
import time import time
while True: while True:
time.sleep(1) time.sleep(1)
except KeyboardInterrupt: except KeyboardInterrupt:
ftp_manager.stop_server() ftp_manager.stop()
print("FTP Server stopped.") print("FTP Server stopped.")
except Exception as e: except Exception as e:
print(f"Error starting FTP server: {e}") print(f"Error starting FTP server: {e}")

View File

@ -63,7 +63,10 @@ class FactExtractionJob:
fact_data = self.ai_processor.extract_facts_from_text(text_content, prompt, model) fact_data = self.ai_processor.extract_facts_from_text(text_content, prompt, model)
# Store fact in database # Store fact in database
self.db_manager.store_fact(fact_data, table_name, file_path) self.db_manager.insert_fact(table_name, {
**fact_data,
'file_path': file_path
})
# Mark file as processed # Mark file as processed
self.db_manager.mark_file_processed(file_path, table_name) self.db_manager.mark_file_processed(file_path, table_name)
@ -127,9 +130,18 @@ class FactExtractionScheduler:
def _run_all_jobs(self): def _run_all_jobs(self):
"""Run all scheduled jobs""" """Run all scheduled jobs"""
print("Running scheduled fact extraction jobs...") print("Running scheduled fact extraction jobs...")
# In a real implementation, this would iterate through configured jobs for job_id, job_info in self.jobs.items():
# For now, we'll just run a basic check try:
pass job = FactExtractionJob(self.config, self.db_manager, self.file_processor, self.ai_processor)
job.execute(
job_info['directory_path'],
job_info['table_name'],
job_info.get('prompt', ''),
job_info.get('model', 'gpt-oss')
)
except Exception as e:
print(f"Error running job {job_id}: {str(e)}")
increment_error()
def add_job(self, directory_path: str, table_name: str, prompt: str = "", def add_job(self, directory_path: str, table_name: str, prompt: str = "",
model: str = "gpt-oss", interval_minutes: int = 10): model: str = "gpt-oss", interval_minutes: int = 10):

View File

@ -1,11 +1,11 @@
flask flask==3.0.0
pyftpdlib pyftpdlib==1.5.6
openai openai==1.12.0
pdfminer.six pdfminer.six==20231228
beautifulsoup4 beautifulsoup4==4.12.3
newspaper3k newspaper3k==0.2.8
apscheduler apscheduler==3.10.4
click click==8.1.7
docker docker==7.1.0
requests requests==2.31.0
prometheus-client prometheus-client==0.20.0