Fix code formatting and linting issues
- Auto-format all Python files with black and isort - Remove unused imports with autoflake - Fix flake8 issues (missing newlines, blank lines, etc.) - Fix nonlocal/global scope issues in main.py - Fix security.py import order (E402) - Remove f-string without placeholders - Add nosec comment for intentional exception handling - Fix test imports to match refactored DMARCParser API Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
+69
-57
@@ -1,19 +1,19 @@
|
||||
from fastapi import FastAPI, Request, BackgroundTasks, Depends
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.templating import Jinja2Templates
|
||||
from fastapi.responses import HTMLResponse
|
||||
import os
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime
|
||||
|
||||
from app.api.api_v1.api import api_router
|
||||
from app.core.config import get_settings
|
||||
from app.core.security import require_admin_auth, generate_api_key, add_api_key
|
||||
from app.core.security import add_api_key, generate_api_key, require_admin_auth
|
||||
from app.middleware.security import SecurityHeadersMiddleware
|
||||
from app.services.imap_client import IMAPClient
|
||||
from app.services.report_store import ReportStore
|
||||
from fastapi import BackgroundTasks, Depends, FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import HTMLResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.templating import Jinja2Templates
|
||||
|
||||
# Set up logging
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -28,51 +28,56 @@ last_check_time = None
|
||||
async def scheduled_imap_polling():
|
||||
"""Background task for periodically checking IMAP for new DMARC reports"""
|
||||
global last_check_time
|
||||
|
||||
|
||||
try:
|
||||
# How often to check for emails (in seconds)
|
||||
check_interval = 3600 # Default: 1 hour
|
||||
|
||||
|
||||
while True:
|
||||
logger.info("Starting scheduled IMAP polling for DMARC reports")
|
||||
|
||||
|
||||
try:
|
||||
# Create IMAP client and fetch reports
|
||||
imap_client = IMAPClient(delete_emails=False)
|
||||
results = imap_client.fetch_reports(days=9999)
|
||||
|
||||
|
||||
# Update last check time
|
||||
last_check_time = datetime.now()
|
||||
|
||||
|
||||
if results["success"]:
|
||||
logger.info(f"IMAP polling completed: {results['processed']} emails processed, "
|
||||
f"{results['reports_found']} reports found")
|
||||
|
||||
logger.info(
|
||||
f"IMAP polling completed: {results['processed']} emails processed, "
|
||||
f"{results['reports_found']} reports found"
|
||||
)
|
||||
|
||||
# If new domains were found, log them
|
||||
if results["new_domains"]:
|
||||
logger.info(f"New domains found: {', '.join(results['new_domains'])}")
|
||||
|
||||
|
||||
else:
|
||||
logger.error(f"IMAP polling failed: {results.get('error', 'Unknown error')}")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in IMAP polling task: {str(e)}")
|
||||
|
||||
|
||||
# Wait for the next check interval
|
||||
await asyncio.sleep(check_interval)
|
||||
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("IMAP polling task cancelled")
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""Create and configure the FastAPI application"""
|
||||
# Task management for background jobs
|
||||
background_task = None
|
||||
|
||||
app = FastAPI(
|
||||
title=settings.PROJECT_NAME,
|
||||
openapi_url=f"{settings.API_V1_STR}/openapi.json",
|
||||
version="0.1.0",
|
||||
)
|
||||
|
||||
|
||||
# Add security headers middleware
|
||||
# Determine environment from settings or environment variable
|
||||
environment = os.getenv("ENVIRONMENT", "development")
|
||||
@@ -88,12 +93,12 @@ def create_app() -> FastAPI:
|
||||
allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
|
||||
# Security: Specify allowed headers instead of wildcard
|
||||
allow_headers=[
|
||||
"Content-Type",
|
||||
"Authorization",
|
||||
"Content-Type",
|
||||
"Authorization",
|
||||
"X-API-Key",
|
||||
"Accept",
|
||||
"Origin",
|
||||
"X-Requested-With"
|
||||
"X-Requested-With",
|
||||
],
|
||||
# Security: Limit exposed headers
|
||||
expose_headers=["Content-Length", "X-RateLimit-Limit"],
|
||||
@@ -102,20 +107,24 @@ def create_app() -> FastAPI:
|
||||
|
||||
# Include API router
|
||||
app.include_router(api_router, prefix=settings.API_V1_STR)
|
||||
|
||||
|
||||
# Mount static files directory
|
||||
app.mount("/static", StaticFiles(directory=os.path.join(os.path.dirname(__file__), "static")), name="static")
|
||||
|
||||
app.mount(
|
||||
"/static",
|
||||
StaticFiles(directory=os.path.join(os.path.dirname(__file__), "static")),
|
||||
name="static",
|
||||
)
|
||||
|
||||
# Set up event handlers for startup and shutdown
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
"""Initialize background tasks and security on application startup"""
|
||||
global background_task
|
||||
|
||||
nonlocal background_task
|
||||
|
||||
# Generate and provide admin API key
|
||||
api_key = generate_api_key()
|
||||
add_api_key(api_key)
|
||||
|
||||
|
||||
# Security: Log only last 8 characters for reference
|
||||
logger.warning(
|
||||
"=" * 80 + "\n"
|
||||
@@ -126,23 +135,23 @@ def create_app() -> FastAPI:
|
||||
"Use this key in the X-API-Key header for admin endpoints.\n"
|
||||
"=" * 80
|
||||
)
|
||||
|
||||
|
||||
# In development, also log the full key for convenience
|
||||
# This should be removed in production or controlled by environment variable
|
||||
if os.getenv("ENVIRONMENT", "development") == "development":
|
||||
logger.info(f"Development Mode - Full API Key: {api_key}")
|
||||
|
||||
|
||||
# Check if IMAP credentials are configured
|
||||
if all([settings.IMAP_SERVER, settings.IMAP_USERNAME, settings.IMAP_PASSWORD]):
|
||||
logger.info("Starting IMAP polling background task")
|
||||
background_task = asyncio.create_task(scheduled_imap_polling())
|
||||
else:
|
||||
logger.warning("IMAP credentials not fully configured, polling disabled")
|
||||
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown_event():
|
||||
"""Clean up background tasks on application shutdown"""
|
||||
global background_task
|
||||
nonlocal background_task # noqa: F824
|
||||
if background_task:
|
||||
logger.info("Cancelling IMAP polling background task")
|
||||
background_task.cancel()
|
||||
@@ -150,7 +159,7 @@ def create_app() -> FastAPI:
|
||||
await background_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
|
||||
return app
|
||||
|
||||
|
||||
@@ -162,7 +171,7 @@ templates = Jinja2Templates(directory=templates_dir)
|
||||
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
async def dashboard(request: Request):
|
||||
async def index(request: Request):
|
||||
return templates.TemplateResponse("index.html", {"request": request})
|
||||
|
||||
|
||||
@@ -173,58 +182,64 @@ async def dashboard(request: Request):
|
||||
"dashboard.html", {"request": request, "app_name": settings.PROJECT_NAME}
|
||||
)
|
||||
|
||||
|
||||
@app.get("/login", response_class=HTMLResponse)
|
||||
async def login(request: Request):
|
||||
return templates.TemplateResponse(
|
||||
"login.html", {"request": request, "app_name": settings.PROJECT_NAME}
|
||||
)
|
||||
|
||||
|
||||
@app.get("/setup", response_class=HTMLResponse)
|
||||
async def setup(request: Request):
|
||||
return templates.TemplateResponse(
|
||||
"setup.html", {"request": request, "app_name": settings.PROJECT_NAME}
|
||||
)
|
||||
|
||||
|
||||
@app.get("/domains", response_class=HTMLResponse)
|
||||
async def domains(request: Request):
|
||||
return templates.TemplateResponse("domains.html", {"request": request})
|
||||
|
||||
|
||||
@app.get("/domain/{domain_id}", response_class=HTMLResponse)
|
||||
async def domain_details(request: Request, domain_id: str):
|
||||
"""View detailed reports for a specific domain"""
|
||||
store = ReportStore.get_instance()
|
||||
domains = store.get_domains()
|
||||
|
||||
|
||||
if domain_id not in domains:
|
||||
# Domain not found, redirect to domains list
|
||||
return templates.TemplateResponse(
|
||||
"domains.html",
|
||||
{"request": request, "error": f"Domain {domain_id} not found"}
|
||||
"domains.html", {"request": request, "error": f"Domain {domain_id} not found"}
|
||||
)
|
||||
|
||||
|
||||
domain_summary = store.get_domain_summary(domain_id)
|
||||
|
||||
|
||||
return templates.TemplateResponse(
|
||||
"domain_details.html",
|
||||
"domain_details.html",
|
||||
{
|
||||
"request": request,
|
||||
"domain_id": domain_id,
|
||||
"domain": {
|
||||
"name": domain_id,
|
||||
"description": "", # Add description if available
|
||||
"policy": domain_summary.get("policy", "unknown")
|
||||
}
|
||||
}
|
||||
"policy": domain_summary.get("policy", "unknown"),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.get("/reports", response_class=HTMLResponse)
|
||||
async def reports(request: Request):
|
||||
return templates.TemplateResponse("reports.html", {"request": request})
|
||||
|
||||
|
||||
@app.get("/settings", response_class=HTMLResponse)
|
||||
async def settings_page(request: Request):
|
||||
return templates.TemplateResponse("settings.html", {"request": request})
|
||||
|
||||
|
||||
@app.get("/upload", response_class=HTMLResponse)
|
||||
async def upload_page(request: Request):
|
||||
return templates.TemplateResponse("upload.html", {"request": request})
|
||||
@@ -233,37 +248,36 @@ async def upload_page(request: Request):
|
||||
# API endpoint to manually trigger IMAP polling
|
||||
@app.post("/api/v1/admin/trigger-poll")
|
||||
async def trigger_imap_poll(
|
||||
background_tasks: BackgroundTasks,
|
||||
auth: dict = Depends(require_admin_auth)
|
||||
background_tasks: BackgroundTasks, auth: dict = Depends(require_admin_auth)
|
||||
):
|
||||
"""
|
||||
Manually trigger IMAP polling (admin only - requires authentication)
|
||||
|
||||
|
||||
Security: Requires either X-API-Key header or Bearer token
|
||||
"""
|
||||
global last_check_time
|
||||
|
||||
|
||||
try:
|
||||
# Create IMAP client and fetch reports
|
||||
imap_client = IMAPClient(delete_emails=False)
|
||||
results = imap_client.fetch_reports(days=7)
|
||||
|
||||
|
||||
# Update last check time
|
||||
last_check_time = datetime.now()
|
||||
|
||||
|
||||
return {
|
||||
"success": results["success"],
|
||||
"timestamp": last_check_time.isoformat(),
|
||||
"processed": results["processed"],
|
||||
"reports_found": results["reports_found"],
|
||||
"new_domains": results["new_domains"],
|
||||
"authenticated_by": auth.get("auth_type")
|
||||
"authenticated_by": auth.get("auth_type"),
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Error triggering IMAP poll: {str(e)}")
|
||||
return {
|
||||
"success": False,
|
||||
"error": "Failed to trigger IMAP poll. Check server logs for details."
|
||||
"error": "Failed to trigger IMAP poll. Check server logs for details.",
|
||||
}
|
||||
|
||||
|
||||
@@ -272,13 +286,11 @@ async def trigger_imap_poll(
|
||||
async def get_poll_status(auth: dict = Depends(require_admin_auth)):
|
||||
"""
|
||||
Get the status of IMAP polling (admin only - requires authentication)
|
||||
|
||||
|
||||
Security: Requires either X-API-Key header or Bearer token
|
||||
"""
|
||||
global last_check_time
|
||||
|
||||
return {
|
||||
"is_running": background_task is not None and not background_task.done(),
|
||||
"last_check": last_check_time.isoformat() if last_check_time else None,
|
||||
"authenticated_by": auth.get("auth_type")
|
||||
}
|
||||
"authenticated_by": auth.get("auth_type"),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user