edb_advanced_lambda.py

import os
import re
import json
import logging
import boto3
import requests
import atoma
import pandas as pd
from datetime import datetime
from typing import List, Dict, Any, Optional
from pathlib import Path
from io import BytesIO

# --- Configuration ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

FEED_URL = os.getenv("FEED_URL", "https://www.exploit-db.com/rss.xml")
S3_BUCKET = os.getenv("S3_BUCKET", "your-bucket-name")
S3_PREFIX = os.getenv("S3_PREFIX", "exploit-db/")
CSV_FILENAME = "EDB.csv"
HTML_FILENAME = "EDB.html"

# Pre-compiled regex
RE_BAD_CHARS = re.compile(r"[\[\]\';\:!\*\\(\)]")

# Initialize S3 client
s3_client = boto3.client('s3')

def clean_string(value: str) -> str:
    if not isinstance(value, str):
        return ""
    return RE_BAD_CHARS.sub("", value).strip()

def extract_data(exploit) -> Dict[str, str]:
    try:
        pub_date = getattr(exploit, 'pub_date', None)
        published = pub_date.strftime('%Y-%m-%d') if pub_date else "Unknown"

        raw_title = getattr(exploit, 'title', '') or ""
        title_parts = raw_title.split(' - ', 1)
        title = title_parts[0].strip() if title_parts else raw_title
        title = clean_string(title)

        link = getattr(exploit, 'link', '') or ""
        link = clean_string(link)

        raw_desc = getattr(exploit, 'description', '') or ""
        
        # Attack Vector: [text]
        vector_match = re.search(r"\[(.*?)\]", raw_desc)
        attack_vector = vector_match.group(1).strip() if vector_match else "N/A"

        # Attack Description: - text (
        desc_match = re.search(r"\s*-\s*(.*?)\s*\(", raw_desc)
        attack_description = desc_match.group(1).strip() if desc_match else "N/A"

        # Attack Type: (text) at end
        type_match = re.search(r"\((.*?)\)\s*$", raw_desc)
        attack_type = type_match.group(1).strip() if type_match else "N/A"

        return {
            "PUBLISHED": published,
            "TITLE": title,
            "LINK": link,
            "ATTACK VECTOR": attack_vector,
            "ATTACK DESCRIPTION": attack_description,
            "ATTACK TYPE": attack_type
        }
    except Exception as e:
        logger.warning(f"Skipping item due to parse error: {e}")
        return {
            "PUBLISHED": "Unknown",
            "TITLE": "Unknown",
            "LINK": "",
            "ATTACK VECTOR": "N/A",
            "ATTACK DESCRIPTION": "N/A",
            "ATTACK TYPE": "N/A"
        }

def fetch_and_process() -> pd.DataFrame:
    try:
        logger.info(f"Fetching feed from {FEED_URL}")
        response = requests.get(FEED_URL, timeout=10)
        response.raise_for_status()
        feed = atoma.parse_rss_bytes(response.content)
        
        logger.info(f"Feed: {feed.title}")
        
        items = []
        for exploit in feed.items:
            items.append(extract_data(exploit))
        
        return pd.DataFrame(items)
    except Exception as e:
        logger.error(f"Failed to fetch or process feed: {e}")
        raise

def save_to_s3(df: pd.DataFrame, filename: str, content_type: str):
    try:
        # Convert DataFrame to string
        if content_type == 'text/html':
            content = df.to_html(index=False)
        else:
            content = df.to_csv(index=False)
        
        # Generate key with timestamp
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        key = f"{S3_PREFIX}{timestamp}_{filename}"
        
        s3_client.put_object(
            Bucket=S3_BUCKET,
            Key=key,
            Body=content.encode('utf-8'),
            ContentType=content_type,
            Metadata={'generated-at': timestamp, 'record-count': str(len(df))}
        )
        
        logger.info(f"Saved {filename} to s3://{S3_BUCKET}/{key}")
        return f"s3://{S3_BUCKET}/{key}"
    except Exception as e:
        logger.error(f"Failed to save to S3: {e}")
        raise

def lambda_handler(event, context):
    """Main Lambda entry point."""
    logger.info(f"Starting execution. Request ID: {context.aws_request_id}")
    
    try:
        # Process data
        df = fetch_and_process()
        
        if df.empty:
            logger.warning("No data processed. Exiting.")
            return {
                'statusCode': 200,
                'body': json.dumps({'message': 'No data found'})
            }
        
        # Save outputs
        csv_url = save_to_s3(df, CSV_FILENAME, 'text/csv')
        html_url = save_to_s3(df, HTML_FILENAME, 'text/html')
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Success',
                'records_processed': len(df),
                'csv_url': csv_url,
                'html_url': html_url
            })
        }
        
    except Exception as e:
        logger.error(f"Execution failed: {e}", exc_info=True)
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }