"""
Website crawler for discovering downloadable files
"""
import os
import re
import time
import hashlib
import requests
from urllib.parse import urljoin, urlparse, parse_qs
from bs4 import BeautifulSoup
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from config import Config
import logging

logger = logging.getLogger(__name__)

class WebsiteCrawler:
    def __init__(self):
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        })
        self.visited_urls = set()
        self.downloadable_files = []
        self.max_depth = Config.MAX_DEPTH
        self.max_pages = Config.MAX_PAGES
        self.crawl_delay = Config.CRAWL_DELAY
        
        # File extensions to look for
        self.file_extensions = {
            '.exe', '.dll', '.msi', '.bat', '.cmd', '.ps1', '.vbs', '.js', '.jar',
            '.zip', '.rar', '.7z', '.tar', '.gz', '.pdf', '.doc', '.docx', '.xls',
            '.xlsx', '.ppt', '.pptx', '.scr', '.com', '.pif', '.lnk'
        }
        
        # Suspicious file patterns
        self.suspicious_patterns = [
            r'keygen', r'crack', r'patch', r'loader', r'injector', r'bypass',
            r'admin', r'root', r'backdoor', r'trojan', r'virus', r'malware'
        ]

    def crawl_website(self, start_url):
        """
        Main crawling function that discovers downloadable files
        """
        logger.info(f"Starting crawl of {start_url}")
        
        try:
            # First, try with requests/BeautifulSoup for static content
            files = self._crawl_static_content(start_url)
            
            # If we need dynamic content, use Selenium
            if len(files) < 5:  # If we didn't find many files, try dynamic crawling
                dynamic_files = self._crawl_dynamic_content(start_url)
                files.extend(dynamic_files)
            
            # Remove duplicates based on URL
            unique_files = []
            seen_urls = set()
            for file_info in files:
                if file_info['url'] not in seen_urls:
                    unique_files.append(file_info)
                    seen_urls.add(file_info['url'])
            
            logger.info(f"Found {len(unique_files)} unique downloadable files")
            return unique_files
            
        except Exception as e:
            logger.error(f"Error crawling website {start_url}: {str(e)}")
            return []

    def _crawl_static_content(self, start_url):
        """Crawl static content using requests and BeautifulSoup"""
        files = []
        urls_to_visit = [(start_url, 0)]  # (url, depth)
        
        while urls_to_visit and len(self.visited_urls) < self.max_pages:
            current_url, depth = urls_to_visit.pop(0)
            
            if current_url in self.visited_urls or depth > self.max_depth:
                continue
                
            self.visited_urls.add(current_url)
            
            try:
                response = self.session.get(current_url, timeout=10)
                response.raise_for_status()
                
                soup = BeautifulSoup(response.content, 'html.parser')
                
                # Find downloadable files
                page_files = self._extract_files_from_page(soup, current_url)
                files.extend(page_files)
                
                # Find links to crawl further
                if depth < self.max_depth:
                    links = self._extract_links(soup, current_url)
                    for link in links:
                        if link not in self.visited_urls:
                            urls_to_visit.append((link, depth + 1))
                
                time.sleep(self.crawl_delay)
                
            except Exception as e:
                logger.warning(f"Error crawling {current_url}: {str(e)}")
                continue
        
        return files

    def _crawl_dynamic_content(self, start_url):
        """Crawl dynamic content using Selenium"""
        files = []
        
        try:
            chrome_options = Options()
            chrome_options.add_argument('--headless')
            chrome_options.add_argument('--no-sandbox')
            chrome_options.add_argument('--disable-dev-shm-usage')
            chrome_options.add_argument('--disable-gpu')
            chrome_options.add_argument('--window-size=1920,1080')
            
            driver = webdriver.Chrome(options=chrome_options)
            driver.get(start_url)
            
            # Wait for page to load
            WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.TAG_NAME, "body"))
            )
            
            # Scroll to load dynamic content
            driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            time.sleep(2)
            
            # Get page source and parse
            soup = BeautifulSoup(driver.page_source, 'html.parser')
            files = self._extract_files_from_page(soup, start_url)
            
            driver.quit()
            
        except Exception as e:
            logger.error(f"Error in dynamic crawling: {str(e)}")
        
        return files

    def _extract_files_from_page(self, soup, base_url):
        """Extract downloadable files from a page"""
        files = []
        
        # Find all links
        links = soup.find_all('a', href=True)
        
        for link in links:
            href = link['href']
            full_url = urljoin(base_url, href)
            
            # Check if it's a downloadable file
            if self._is_downloadable_file(full_url, link.get_text()):
                file_info = {
                    'url': full_url,
                    'filename': self._extract_filename(full_url, link.get_text()),
                    'path': href,
                    'page_url': base_url,
                    'link_text': link.get_text().strip(),
                    'suspicious': self._is_suspicious_filename(full_url, link.get_text())
                }
                files.append(file_info)
        
        # Also check for embedded files in scripts, iframes, etc.
        embedded_files = self._find_embedded_files(soup, base_url)
        files.extend(embedded_files)
        
        return files

    def _extract_links(self, soup, base_url):
        """Extract internal links for further crawling"""
        links = []
        
        for link in soup.find_all('a', href=True):
            href = link['href']
            full_url = urljoin(base_url, href)
            
            # Only crawl internal links
            if self._is_internal_link(full_url, base_url):
                links.append(full_url)
        
        return links

    def _is_downloadable_file(self, url, link_text):
        """Check if URL points to a downloadable file"""
        parsed_url = urlparse(url)
        path = parsed_url.path.lower()
        
        # Check file extension
        for ext in self.file_extensions:
            if path.endswith(ext):
                return True
        
        # Check for download indicators in URL
        download_indicators = ['download', 'file', 'attachment']
        for indicator in download_indicators:
            if indicator in path or indicator in url.lower():
                return True
        
        # Check link text for download indicators
        if link_text:
            link_text_lower = link_text.lower()
            download_texts = ['download', 'click here', 'get file', 'save']
            for text in download_texts:
                if text in link_text_lower:
                    return True
        
        return False

    def _is_suspicious_filename(self, url, link_text):
        """Check if filename or link text contains suspicious patterns"""
        text_to_check = f"{url} {link_text}".lower()
        
        for pattern in self.suspicious_patterns:
            if re.search(pattern, text_to_check):
                return True
        
        return False

    def _extract_filename(self, url, link_text):
        """Extract filename from URL or link text"""
        parsed_url = urlparse(url)
        path = parsed_url.path
        
        # Try to get filename from URL path
        if path and '/' in path:
            filename = path.split('/')[-1]
            if filename and '.' in filename:
                return filename
        
        # Try to get filename from link text
        if link_text and link_text.strip():
            # Clean up link text
            filename = re.sub(r'[^\w\-_\.]', '_', link_text.strip())
            if filename:
                return filename
        
        # Generate filename from URL
        return f"file_{hashlib.md5(url.encode()).hexdigest()[:8]}"

    def _is_internal_link(self, url, base_url):
        """Check if URL is internal to the same domain"""
        try:
            base_domain = urlparse(base_url).netloc
            url_domain = urlparse(url).netloc
            return base_domain == url_domain
        except:
            return False

    def _find_embedded_files(self, soup, base_url):
        """Find files embedded in scripts, iframes, etc."""
        files = []
        
        # Check script tags for file references
        scripts = soup.find_all('script')
        for script in scripts:
            if script.string:
                # Look for file URLs in JavaScript
                file_urls = re.findall(r'["\']([^"\']*\.(?:exe|dll|zip|pdf|doc))["\']', script.string)
                for url in file_urls:
                    full_url = urljoin(base_url, url)
                    if self._is_downloadable_file(full_url, ''):
                        file_info = {
                            'url': full_url,
                            'filename': self._extract_filename(full_url, ''),
                            'path': url,
                            'page_url': base_url,
                            'link_text': 'Embedded in script',
                            'suspicious': self._is_suspicious_filename(full_url, '')
                        }
                        files.append(file_info)
        
        return files

    def download_file(self, file_info, download_dir='/tmp'):
        """Download a file for analysis"""
        try:
            response = self.session.get(file_info['url'], stream=True, timeout=30)
            response.raise_for_status()
            
            # Create download directory if it doesn't exist
            os.makedirs(download_dir, exist_ok=True)
            
            # Generate safe filename
            safe_filename = re.sub(r'[^\w\-_\.]', '_', file_info['filename'])
            file_path = os.path.join(download_dir, safe_filename)
            
            # Download file
            with open(file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            
            # Get file info
            file_size = os.path.getsize(file_path)
            file_hash = self._calculate_file_hash(file_path)
            
            file_info.update({
                'local_path': file_path,
                'file_size': file_size,
                'file_hash': file_hash,
                'content_type': response.headers.get('content-type', '')
            })
            
            return file_info
            
        except Exception as e:
            logger.error(f"Error downloading file {file_info['url']}: {str(e)}")
            return None

    def _calculate_file_hash(self, file_path):
        """Calculate SHA256 hash of file"""
        sha256_hash = hashlib.sha256()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                sha256_hash.update(chunk)
        return sha256_hash.hexdigest()
