backend/models.py

import sys
import os
sys.path.insert(0, os.path.abspath('backend'))

from bs4 import BeautifulSoup
from celery import Celery
from celery.schedules import crontab
from collections import defaultdict
from scanners.base_scanner import (
    active_scan_logs, add_log, get_scan_logs, parse_domain,
    cleanup_scan_logs, schedule_log_cleanup, emit_scan_progress
)
from scanners import get_pipeline, get_phases, build_scanner, apply_scan_options
try:
    from backend.utils.fuzzer_engine import ContextAwareFuzzer
except ImportError:
    from utils.fuzzer_engine import ContextAwareFuzzer
from cryptography import x509
from cryptography.hazmat.backends import default_backend
from datetime import datetime, timezone
from datetime import datetime, timezone, timedelta
from datetime import datetime, timezone, timezone
from dotenv import load_dotenv
load_dotenv()

import stripe
from flask import Blueprint, request, jsonify, current_app, send_from_directory
from werkzeug.utils import secure_filename
from flask import Blueprint, send_file, jsonify, request
from flask import Flask
from flask import jsonify
from flask import render_template
from flask import request, abort, g, Response, make_response
from flask_cors import CORS
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from flask_socketio import SocketIO, emit, join_room, leave_room
from flask_sqlalchemy import SQLAlchemy
from functools import wraps
from markupsafe import escape  # always available with Flask
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfgen import canvas
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak, Image, Flowable
from reportlab.graphics.shapes import Drawing
from reportlab.graphics.charts.barcharts import VerticalBarChart
from sqlalchemy import event
from sqlalchemy import func
from sqlalchemy import inspect, text
from sqlalchemy import text
from sqlalchemy.engine import Engine
from typing import Any
from typing import Any, Callable
from typing import Callable
from typing import Literal
from urllib.parse import urljoin, urlparse
from urllib.parse import urlparse
import base64
import bcrypt
import concurrent.futures
from backend.utils.email_service import (
    send_welcome_email, 
    send_scan_started, 
    send_scan_completed, 
    send_scan_failed,
    send_critical_alert
)

import hashlib
import html
import io
import itertools
import json
import jwt
import math
import os
import re
import re, time, ipaddress, os, hashlib, threading
import requests
import socket
import sqlite3
import ssl
import statistics
import threading
import time
import traceback
import urllib.error
import urllib.parse
import urllib.request
import urllib3
import uuid
import ipaddress



try:
    from .extensions import db, celery, socketio, limiter
except ImportError:
    from extensions import db, celery, socketio, limiter


# --- From models.py ---

class Organization(db.Model):
    __tablename__ = 'organizations'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    name = db.Column(db.String(150), nullable=False)
    subscription_tier = db.Column(db.String(50), nullable=False, default='free')
    status = db.Column(db.String(50), default='active')
    api_key = db.Column(db.String(100), unique=True, nullable=True)
    webhook_url = db.Column(db.String(500), nullable=True)
    report_logo_url = db.Column(db.String(500), nullable=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    users = db.relationship('User', backref='organization', lazy=True)
    scans = db.relationship('Scan', backref='organization', lazy=True, cascade="all, delete-orphan")


class Role(db.Model):
    __tablename__ = 'roles'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    name = db.Column(db.String(50), unique=True, nullable=False)

class User(db.Model):
    __tablename__ = 'users'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    email = db.Column(db.String(120), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    first_name = db.Column(db.String(100), nullable=True)
    last_name = db.Column(db.String(100), nullable=True)
    role = db.Column(db.String(50), nullable=False, default='org_admin') # super_admin, support_engineer, org_admin, soc_analyst, executive, read_only
    org_id = db.Column(db.String(36), db.ForeignKey('organizations.id'), nullable=True)
    
    failed_login_attempts = db.Column(db.Integer, default=0)
    locked_until = db.Column(db.DateTime, nullable=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    # Policy Consent & Security Fields
    terms_accepted_at = db.Column(db.DateTime, nullable=True)
    privacy_policy_agreed_at = db.Column(db.DateTime, nullable=True)
    policy_version_agreed = db.Column(db.String(20), default='v1.0')
    mfa_enabled = db.Column(db.Boolean, default=False)
    mfa_secret = db.Column(db.String(100), nullable=True)
    reset_token = db.Column(db.String(255), nullable=True)
    reset_token_expires = db.Column(db.DateTime, nullable=True)
    email_verified = db.Column(db.Boolean, default=False)
    
    scans = db.relationship('Scan', backref='user', lazy=True, cascade="all, delete-orphan")
    alert_settings = db.relationship('AlertSettings', backref='user', uselist=False, lazy=True, cascade="all, delete-orphan")
    
    @property
    def subscription_tier(self):
        if self.organization:
            return self.organization.subscription_tier
        return 'free'

    @subscription_tier.setter
    def subscription_tier(self, value):
        if self.organization:
            self.organization.subscription_tier = value
    
    def set_password(self, password):
        salt = bcrypt.gensalt()
        self.password_hash = bcrypt.hashpw(password.encode('utf-8'), salt).decode('utf-8')
        
    def check_password(self, password):
        if not self.password_hash or not password:
            return False
        try:
            return bcrypt.checkpw(password.encode('utf-8'), self.password_hash.encode('utf-8'))
        except (ValueError, TypeError, AttributeError):
            return False

class SubscriptionTier(db.Model):
    __tablename__ = 'subscription_tiers'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    id = db.Column(db.String(50), primary_key=True) 
    name = db.Column(db.String(100), nullable=False)
    monthly_price = db.Column(db.Integer, nullable=False, default=0) 
    yearly_price = db.Column(db.Integer, nullable=False, default=0)

class AuditLog(db.Model):
    __tablename__ = 'audit_logs'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    admin_id = db.Column(db.String(36), nullable=False)
    action = db.Column(db.String(255), nullable=False)
    target_id = db.Column(db.String(100), nullable=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

class Payment(db.Model):
    __tablename__ = 'payments'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    org_id = db.Column(db.String(36), db.ForeignKey('organizations.id', ondelete='CASCADE'), nullable=True)
    user_id = db.Column(db.String(36), db.ForeignKey('users.id', ondelete='CASCADE'), nullable=False)
    razorpay_payment_id = db.Column(db.String(100), unique=True, nullable=True)
    razorpay_order_id = db.Column(db.String(100), nullable=True)
    stripe_session_id = db.Column(db.String(100), unique=True, nullable=True)
    stripe_payment_id = db.Column(db.String(100), unique=True, nullable=True)
    tier_id = db.Column(db.String(50), nullable=False)
    amount = db.Column(db.Integer, nullable=False)
    currency = db.Column(db.String(10), default='USD')
    status = db.Column(db.String(50), default='successful')
    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))

class Scan(db.Model):
    __tablename__ = 'scans'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    org_id = db.Column(db.String(36), db.ForeignKey('organizations.id'), nullable=False)
    user_id = db.Column(db.String(36), db.ForeignKey('users.id'), nullable=False)
    target_url = db.Column(db.String(500), nullable=False)
    scan_type = db.Column(db.String(50), nullable=False, default='Full') # Full, Port, SSL, OWASP
    status = db.Column(db.String(50), nullable=False, default='queued') # queued, scanning, completed, failed
    security_score = db.Column(db.Integer, nullable=True) # 0 to 100
    started_at = db.Column(db.DateTime, default=datetime.utcnow)
    completed_at = db.Column(db.DateTime, nullable=True)
    auth_headers = db.Column(db.JSON, nullable=True) # For authenticated scanning
    scan_options = db.Column(db.JSON, nullable=True) # crawl_depth, exclude_paths, enable_red_team
    ssl_info = db.Column(db.JSON, nullable=True) # Cached SSL certificate info
    
    duration_seconds = db.Column(db.Integer, nullable=True)
    error_message = db.Column(db.Text, nullable=True)
    critical_count = db.Column(db.Integer, default=0)
    high_count = db.Column(db.Integer, default=0)
    medium_count = db.Column(db.Integer, default=0)
    low_count = db.Column(db.Integer, default=0)
    
    vulnerabilities = db.relationship('Vulnerability', backref='scan', lazy=True, cascade="all, delete-orphan")

    def __init__(self, org_id, user_id, target_url, scan_type='Full', status='queued', security_score=None, started_at=None, completed_at=None, auth_headers=None, scan_options=None, id=None):
        if id: self.id = id
        self.org_id = org_id
        self.user_id = user_id
        self.target_url = target_url
        self.scan_type = scan_type
        self.status = status
        self.security_score = security_score
        self.auth_headers = auth_headers
        self.scan_options = scan_options
        if started_at: self.started_at = started_at
        if completed_at: self.completed_at = completed_at

class Vulnerability(db.Model):
    __tablename__ = 'vulnerabilities'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id = db.Column(db.String(36), db.ForeignKey('scans.id'), nullable=False)
    title = db.Column(db.String(200), nullable=False)
    severity = db.Column(db.String(50), nullable=False) # Critical, High, Medium, Low
    category = db.Column(db.String(100), nullable=False) # SSL/TLS, Port, Security Headers, Injection
    description = db.Column(db.Text, nullable=False)
    remediation = db.Column(db.Text, nullable=False)
    cvss_score = db.Column(db.Float, nullable=False)
    detected_at = db.Column(db.DateTime, default=datetime.utcnow)
    evidence = db.Column(db.Text, default="")
    payload = db.Column(db.Text, default="")
    request_details = db.Column(db.Text, default="")
    response_details = db.Column(db.Text, default="")
    is_false_positive = db.Column(db.Boolean, default=False)
    cwe_ids = db.Column(db.JSON, nullable=True)
    owasp_category = db.Column(db.String(100), nullable=True)
    exploit_poc = db.Column(db.JSON, nullable=True)
    remediation_code = db.Column(db.Text, nullable=True)
    
    status = db.Column(db.String(50), nullable=False, default='open') # open, in_progress, resolved, ignored
    remediated_at = db.Column(db.DateTime, nullable=True)
    remediated_by = db.Column(db.String(36), nullable=True)

    def __init__(self, scan_id, title, severity, category, description, remediation, cvss_score, detected_at=None, id=None, evidence="", payload="", request_details="", response_details="", cwe_ids=None, owasp_category=None, exploit_poc=None, remediation_code=None):
        if id: self.id = id
        self.scan_id = scan_id
        self.title = title
        self.severity = severity
        self.category = category
        self.description = description
        self.remediation = remediation
        self.cvss_score = cvss_score
        if detected_at: self.detected_at = detected_at
        self.evidence = evidence
        self.payload = payload
        self.request_details = request_details
        self.response_details = response_details
        self.cwe_ids = cwe_ids
        self.owasp_category = owasp_category
        self.exploit_poc = exploit_poc
        self.remediation_code = remediation_code

class DemoBooking(db.Model):
    __tablename__ = 'demo_bookings'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    email = db.Column(db.String(255), nullable=False)
    company_size = db.Column(db.String(100), nullable=False)
    meeting_date = db.Column(db.String(100), nullable=False)
    meeting_time = db.Column(db.String(50), nullable=False)
    status = db.Column(db.String(50), nullable=False, default='pending')
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

    def __init__(self, email, company_size, meeting_date, meeting_time, status='pending', id=None):
        if id: self.id = id
        self.email = email
        self.company_size = company_size
        self.meeting_date = meeting_date
        self.meeting_time = meeting_time
        self.status = status

    def to_dict(self):
        return {
            'id': self.id,
            'email': self.email,
            'company_size': self.company_size,
            'meeting_date': self.meeting_date,
            'meeting_time': self.meeting_time,
            'status': self.status,
            'created_at': self.created_at.isoformat() + 'Z' if self.created_at else None
        }

class EmailLog(db.Model):
    __tablename__ = 'email_logs'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    recipient = db.Column(db.String(255), nullable=False)
    subject = db.Column(db.String(255), nullable=False)
    status = db.Column(db.String(50), nullable=False, default='sent')
    error_message = db.Column(db.Text, nullable=True)
    sent_at = db.Column(db.DateTime, default=datetime.utcnow)

    def __init__(self, recipient, subject, status='sent', error_message=None, id=None):
        if id: self.id = id
        self.recipient = recipient
        self.subject = subject
        self.status = status
        self.error_message = error_message

class ScheduledScan(db.Model):
    __tablename__ = 'scheduled_scans'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    org_id = db.Column(db.String(36), db.ForeignKey('organizations.id'), nullable=False)
    user_id = db.Column(db.String(36), db.ForeignKey('users.id'), nullable=False)
    target_url = db.Column(db.String(500), nullable=False)
    scan_type = db.Column(db.String(50), nullable=False, default='Full')
    frequency = db.Column(db.String(50), nullable=False, default='daily') # daily, weekly, monthly, once
    schedule_time = db.Column(db.String(5), nullable=True) # HH:MM format like "20:00"
    day_of_week = db.Column(db.String(20), nullable=True) # monday, tuesday, etc.
    day_of_month = db.Column(db.Integer, nullable=True) # 1..31
    specific_date = db.Column(db.String(20), nullable=True) # YYYY-MM-DD
    is_active = db.Column(db.Boolean, default=True)
    auth_headers = db.Column(db.JSON, nullable=True) # For authenticated scanning
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    last_run_at = db.Column(db.DateTime, nullable=True)

    def __init__(self, org_id, user_id, target_url, scan_type='Full', frequency='daily', schedule_time=None, day_of_week=None, day_of_month=None, specific_date=None, is_active=True, last_run_at=None, auth_headers=None, id=None):
        if id: self.id = id
        self.org_id = org_id
        self.user_id = user_id
        self.target_url = target_url
        self.scan_type = scan_type
        self.frequency = frequency
        self.schedule_time = schedule_time
        self.day_of_week = day_of_week
        self.day_of_month = day_of_month
        self.specific_date = specific_date
        self.is_active = is_active
        self.auth_headers = auth_headers
        if last_run_at: self.last_run_at = last_run_at

class AlertSettings(db.Model):
    __tablename__ = 'alert_settings'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    user_id = db.Column(db.String(36), db.ForeignKey('users.id'), nullable=False)
    email_notifications = db.Column(db.Boolean, default=True)
    webhook_url = db.Column(db.String(500), nullable=True)
    severity_threshold = db.Column(db.String(50), default='Medium') # Low, Medium, High, Critical


class OrganizationScanQuota(db.Model):
    __tablename__ = 'organization_scan_quotas'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    org_id = db.Column(db.String(36), db.ForeignKey('organizations.id'), nullable=False)
    scan_type = db.Column(db.String(50), nullable=False)
    allocated_count = db.Column(db.Integer, nullable=False, default=0)
    used_count = db.Column(db.Integer, nullable=False, default=0)
    
    __table_args__ = (db.UniqueConstraint('org_id', 'scan_type', name='uix_org_scan_type'),)


class Report(db.Model):
    __tablename__ = 'reports'
    def __init__(self, **kwargs): super().__init__(**kwargs)
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id = db.Column(db.String(36), db.ForeignKey('scans.id', ondelete='CASCADE'), nullable=False)
    org_id = db.Column(db.String(36), db.ForeignKey('organizations.id', ondelete='CASCADE'), nullable=False)
    report_type = db.Column(db.String(50), nullable=False) # pdf, executive_summary, compliance
    file_path = db.Column(db.String(500), nullable=True)
    generated_at = db.Column(db.DateTime, default=datetime.utcnow)