batch_scanner.py

"""
Batch URL Scanner - Process multiple URLs in parallel
"""
import asyncio
from typing import List, Dict, Any
from loguru import logger
from datetime import datetime

class BatchScanner:
    """Handle batch scanning of multiple URLs"""
    
    def __init__(self, aggregator, max_parallel: int = 3):
        self.aggregator = aggregator
        self.max_parallel = max_parallel
        self.queue = asyncio.Queue()
        self.results = []
        self.total = 0
        self.completed = 0
        
    async def scan_batch(self, urls: List[str], progress_callback=None, cancel_check=None) -> List[Dict[str, Any]]:
        """Scan multiple URLs with progress tracking"""
        self.total = len(urls)
        self.completed = 0
        self.results = []
        self.cancel_check = cancel_check
        
        # Add URLs to queue
        for url in urls:
            await self.queue.put(url)
        
        # Create worker tasks
        workers = [
            asyncio.create_task(self._worker(i, progress_callback))
            for i in range(min(self.max_parallel, len(urls)))
        ]
        
        # Wait for all workers to finish
        await self.queue.join()
        
        # Cancel workers
        for worker in workers:
            worker.cancel()
        
        return self.results
    
    async def _worker(self, worker_id: int, progress_callback):
        """Worker task to process URLs from queue"""
        while True:
            try:
                url = await self.queue.get()
                
                # Check for cancel
                if self.cancel_check and self.cancel_check():
                    self.queue.task_done()
                    break
                
                logger.info(f"Worker {worker_id} scanning: {url}")
                
                # Scan URL
                result = await self.aggregator.scan_url(url)
                self.results.append(result)
                
                self.completed += 1
                
                # Call progress callback
                if progress_callback:
                    await progress_callback(self.completed, self.total, url)
                
                self.queue.task_done()
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Worker {worker_id} error: {e}")
                self.queue.task_done()
    
    def get_progress(self) -> Dict[str, Any]:
        """Get current progress"""
        return {
            'total': self.total,
            'completed': self.completed,
            'percentage': (self.completed / self.total * 100) if self.total > 0 else 0,
            'remaining': self.total - self.completed
        }

class DeepScanner:
    """Deep scan a URL by crawling linked pages"""
    
    def __init__(self, aggregator, max_depth: int = 2, max_pages: int = 10):
        self.aggregator = aggregator
        self.max_depth = max_depth
        self.max_pages = max_pages
        
    async def deep_scan(self, url: str, progress_callback=None, cancel_check=None) -> Dict[str, Any]:
        """Perform deep scan with link crawling"""
        from urllib.parse import urljoin, urlparse
        import httpx
        from bs4 import BeautifulSoup
        
        # Normalize URL to ensure it has protocol
        if not url.startswith(('http://', 'https://')):
            url = 'https://' + url
        
        visited = set()
        to_visit = [(url, 0)]  # (url, depth)
        results = []
        
        base_domain = urlparse(url).netloc
        
        while to_visit and len(visited) < self.max_pages:
            # Check for cancellation
            if cancel_check and cancel_check():
                logger.info("Deep scan cancelled by user")
                break
                
            current_url, depth = to_visit.pop(0)
            
            if current_url in visited or depth > self.max_depth:
                continue
            
            visited.add(current_url)
            
            # Scan current URL
            logger.info(f"Deep scanning: {current_url} (depth: {depth})")
            result = await self.aggregator.scan_url(current_url)
            results.append(result)
            
            if progress_callback:
                await progress_callback(len(visited), self.max_pages, current_url)
            
            # Extract links if not at max depth
            if depth < self.max_depth:
                try:
                    logger.info(f"Extracting links from {current_url} (depth {depth}/{self.max_depth})")
                    
                    # Try with different User-Agent strategies
                    headers_options = [
                        {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'},
                        {'User-Agent': 'Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)'},
                        {}  # No User-Agent
                    ]
                    
                    response = None
                    for headers in headers_options:
                        try:
                            async with httpx.AsyncClient(
                                timeout=10.0,
                                follow_redirects=True,
                                headers=headers
                            ) as client:
                                response = await client.get(current_url)
                                if response.status_code == 200:
                                    break
                        except Exception as e:
                            logger.debug(f"Request failed with headers {headers}: {e}")
                            continue
                    
                    if not response:
                        logger.warning(f"All request attempts failed for {current_url}")
                        continue
                        
                    logger.info(f"Got response: {response.status_code}, content-type: {response.headers.get('content-type', 'unknown')}")
                    
                    if response.status_code == 200 and 'text/html' in response.headers.get('content-type', ''):
                        soup = BeautifulSoup(response.text, 'html.parser')
                        all_links = soup.find_all('a', href=True)
                        logger.info(f"Found {len(all_links)} total <a> tags on page")
                        
                        # Extract all links
                        links_found = 0
                        for link in all_links:
                            href = link['href']
                            
                            # Skip non-http(s) links
                            if href.startswith(('#', 'javascript:', 'mailto:', 'tel:')):
                                continue
                            
                            # Convert to absolute URL
                            absolute_url = urljoin(current_url, href)
                            parsed = urlparse(absolute_url)
                            
                            # Only crawl same domain, http/https only
                            if (parsed.scheme in ['http', 'https'] and 
                                parsed.netloc == base_domain and 
                                absolute_url not in visited and
                                absolute_url not in [u for u, d in to_visit]):
                                to_visit.append((absolute_url, depth + 1))
                                links_found += 1
                                logger.debug(f"Added link: {absolute_url}")
                                
                                # Stop if we have enough links queued
                                if len(to_visit) >= self.max_pages:
                                    break
                        
                        logger.info(f"Found {links_found} same-domain links on {current_url}, queue now has {len(to_visit)} URLs")
                    elif response.status_code == 403:
                        logger.warning(f"403 Forbidden - Site blocking bot. Skipping link extraction for {current_url}")
                    else:
                        logger.warning(f"Cannot extract links: status={response.status_code}, content-type={response.headers.get('content-type')}")
                                    
                except Exception as e:
                    logger.warning(f"Error extracting links from {current_url}: {e}")
        
        return {
            'base_url': url,
            'pages_scanned': len(visited),
            'max_depth_reached': max([d for _, d in [(u, 0) for u in visited]]) if visited else 0,
            'results': results,
            'visited_urls': list(visited)
        }

class ScanQueue:
    """Queue system for managing multiple scan requests"""
    
    def __init__(self, max_concurrent: int = 2):
        self.queue = asyncio.Queue()
        self.max_concurrent = max_concurrent
        self.active_scans = {}
        self.completed_scans = {}
        
    async def add_to_queue(self, user_id: int, url: str, scan_type: str = 'standard'):
        """Add scan to queue"""
        scan_id = f"{user_id}_{len(self.active_scans) + len(self.completed_scans)}"
        
        scan_job = {
            'id': scan_id,
            'user_id': user_id,
            'url': url,
            'type': scan_type,
            'status': 'queued',
            'position': self.queue.qsize() + 1,
            'added_at': datetime.now()
        }
        
        await self.queue.put(scan_job)
        return scan_id
    
    async def process_queue(self, aggregator, callback):
        """Process queue with concurrent workers"""
        workers = [
            asyncio.create_task(self._queue_worker(aggregator, callback))
            for _ in range(self.max_concurrent)
        ]
        
        await asyncio.gather(*workers, return_exceptions=True)
    
    async def _queue_worker(self, aggregator, callback):
        """Worker to process queued scans"""
        while True:
            try:
                job = await self.queue.get()
                
                job['status'] = 'processing'
                self.active_scans[job['id']] = job
                
                # Perform scan
                result = await aggregator.scan_url(job['url'])
                
                # Mark complete
                job['status'] = 'completed'
                job['result'] = result
                job['completed_at'] = datetime.now()
                
                self.completed_scans[job['id']] = job
                del self.active_scans[job['id']]
                
                # Callback to notify user
                await callback(job)
                
                self.queue.task_done()
                
            except Exception as e:
                logger.error(f"Queue worker error: {e}")
                if 'job' in locals():
                    job['status'] = 'failed'
                    job['error'] = str(e)
    
    def get_position(self, scan_id: str) -> int:
        """Get position in queue"""
        # Simple implementation - could be improved
        return self.queue.qsize()
    
    def get_status(self, scan_id: str) -> Dict[str, Any]:
        """Get scan status"""
        if scan_id in self.active_scans:
            return self.active_scans[scan_id]
        elif scan_id in self.completed_scans:
            return self.completed_scans[scan_id]
        else:
            return {'status': 'not_found'}