ADD export function to tables

main.py

@@ -3,7 +3,7 @@ import re
 import io
 import datetime
 from fastapi import FastAPI, Request, HTTPException, Query, UploadFile, File, Form
-from fastapi.responses import HTMLResponse, RedirectResponse
+from fastapi.responses import HTMLResponse, RedirectResponse, Response
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 import uvicorn
@@ -350,6 +350,60 @@ def delete_upload(upload_id):
         print(f"Error deleting upload {upload_id}: {e}")
         return False
 
+# CSV Export Functions
+def domains_to_csv(domains: List[Dict]) -> str:
+    """Convert domains data to CSV format"""
+    csv_output = io.StringIO()
+    
+    if not domains:
+        return ""
+        
+    # Determine fields based on data
+    # Always include the full_domain field
+    fields = ["full_domain", "timestamp"]
+    if "base_domain" in domains[0]:
+        fields.insert(1, "base_domain")
+    
+    # Add headers
+    writer = csv.DictWriter(csv_output, fieldnames=fields, extrasaction='ignore')
+    writer.writeheader()
+    
+    # Add data
+    for domain in domains:
+        # Create a row dict with formatted timestamp
+        row = {k: domain.get(k) for k in fields}
+        if "timestamp" in row and row["timestamp"]:
+            # Format timestamp nicely for CSV
+            row["timestamp"] = row["timestamp"].replace('T', ' ').split('.')[0]
+        writer.writerow(row)
+    
+    return csv_output.getvalue()
+
+def dns_records_to_csv(records: List[Dict]) -> str:
+    """Convert DNS records data to CSV format"""
+    csv_output = io.StringIO()
+    
+    if not records:
+        return ""
+    
+    # Define the fields to include in the CSV
+    fields = ["domain", "ttl", "record_class", "record_type", "record_data", "timestamp"]
+    
+    # Add headers
+    writer = csv.DictWriter(csv_output, fieldnames=fields, extrasaction='ignore')
+    writer.writeheader()
+    
+    # Add data
+    for record in records:
+        # Create a row dict with formatted timestamp
+        row = {k: record.get(k) for k in fields}
+        if "timestamp" in row and row["timestamp"]:
+            # Format timestamp nicely for CSV
+            row["timestamp"] = row["timestamp"].replace('T', ' ').split('.')[0]
+        writer.writerow(row)
+    
+    return csv_output.getvalue()
+
 # Routes
 @app.get("/", response_class=HTMLResponse)
 async def home(
@@ -471,6 +525,57 @@ async def dns_records(
         }
     )
 
+@app.get("/export-domains-csv")
+async def export_domains_csv(
+    upload_id: Optional[str] = None,
+    base_domains_only: Optional[bool] = False
+):
+    """Export domains as CSV"""
+    domains = load_domains(upload_id, base_domains_only)
+    csv_content = domains_to_csv(domains)
+    
+    # Generate a filename with timestamp
+    filename = f"domains_export_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
+    
+    # Return the CSV as a downloadable file
+    return Response(
+        content=csv_content,
+        media_type="text/csv",
+        headers={"Content-Disposition": f"attachment; filename={filename}"}
+    )
+
+@app.get("/export-dns-csv")
+async def export_dns_csv(
+    upload_id: Optional[str] = None,
+    record_type: Optional[str] = None,
+    record_class: Optional[str] = None,
+    domain: Optional[str] = None,
+    deduplicate: Optional[bool] = True
+):
+    """Export DNS records as CSV"""
+    # Get entries with applied filters
+    entries = load_dns_entries(upload_id, deduplicate)
+    
+    # Apply additional filters if provided
+    if record_type:
+        entries = [e for e in entries if e.get("record_type") == record_type]
+    if record_class:
+        entries = [e for e in entries if e.get("record_class") == record_class]
+    if domain:
+        entries = [e for e in entries if domain.lower() in e.get("domain", "").lower()]
+    
+    csv_content = dns_records_to_csv(entries)
+    
+    # Generate a filename with timestamp
+    filename = f"dns_records_export_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
+    
+    # Return the CSV as a downloadable file
+    return Response(
+        content=csv_content,
+        media_type="text/csv",
+        headers={"Content-Disposition": f"attachment; filename={filename}"}
+    )
+
 # API Routes
 @app.get("/api/uploads", response_model=List[Dict])
 async def get_all_uploads():