domain-scraping-utils/bin/alive-check.py

#!/usr/bin/env python
import os
import csv
import aiohttp
import asyncio
import ssl
import certifi
import time
import logging
import socket
from aiohttp import ClientSession, TCPConnector, ClientTimeout
from yarl import URL

# Global variables for progress tracking
total_domains = 0
processed_domains = 0

# Custom logging formatter
class ProgressFormatter(logging.Formatter):
    def format(self, record):
        global total_domains, processed_domains
        progress = (processed_domains / total_domains) * 100 if total_domains > 0 else 0
        record.progress = f"[{progress:.2f}% - {processed_domains}/{total_domains}]"
        return super().format(record)

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
handler = logging.StreamHandler()
handler.setFormatter(ProgressFormatter('%(progress)s : %(message)s'))
logger.handlers.clear()
logger.addHandler(handler)

# Determine project root dynamically
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(SCRIPT_DIR)

INPUT_FILE = os.path.join(PROJECT_ROOT, "01-normalized-data.csv")
OUTPUT_FILE = os.path.join(PROJECT_ROOT, "02-alive-check.csv")

MAX_CONCURRENT = 100
CHUNK_SIZE = 500

# Timeout settings
CONNECT_TIMEOUT = 10  # seconds
TOTAL_TIMEOUT = 30  # seconds

# Retry settings
CHUNK_RETRIES = 3
CHUNK_RETRY_DELAY = 60  # seconds

# User agent string
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"

# Status codes that should result in immediate "no"
NO_RETRY_STATUS_CODES = {403}

# Create a SSL context that ignores certificate validation
ssl_context = ssl.create_default_context(cafile=certifi.where())
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE

class CSVManager:
    def __init__(self, filename, mode='w'):
        self.filename = filename
        self.file = open(filename, mode, newline='', encoding='utf-8')
        self.writer = csv.writer(self.file)
        self.lock = asyncio.Lock()
        
        # Write header if file is empty
        if self.file.tell() == 0:
            self.writer.writerow(["domain", "alive", "website"])
            self.file.flush()

    async def write_row(self, row):
        async with self.lock:
            self.writer.writerow(row)
            self.file.flush()

    def close(self):
        self.file.close()

async def check_website(session, domain, csv_manager):
    global processed_domains
    await csv_manager.write_row([domain, "checking", ""])
    logging.info(f"Checking: {domain}")

    for protocol in ['https://', 'http://']:
        url = URL(f"{protocol}{domain}")
        try:
            timeout = ClientTimeout(total=TOTAL_TIMEOUT, connect=CONNECT_TIMEOUT)
            async with session.get(url, timeout=timeout, allow_redirects=True, ssl=ssl_context) as response:
                if response.status < 400:
                    final_url = str(response.url)
                    await csv_manager.write_row([domain, "yes", final_url])
                    logging.info(f"Success: {domain} -> {final_url}")
                    processed_domains += 1
                    return
                elif response.status in NO_RETRY_STATUS_CODES:
                    await csv_manager.write_row([domain, "no", ""])
                    logging.info(f"No retry status code {response.status}: {domain}")
                    processed_domains += 1
                    return
                else:
                    logging.info(f"HTTP Error: {url} - Status {response.status}")
        except asyncio.TimeoutError:
            logging.info(f"Timeout: {url}")
        except Exception as e:
            logging.info(f"Error: {url} - {str(e)}")

    await csv_manager.write_row([domain, "no", ""])
    logging.info(f"Failed: {domain} - Both HTTPS and HTTP attempts failed")
    processed_domains += 1

async def process_chunk(domains, csv_manager, chunk_number):
    for retry in range(CHUNK_RETRIES):
        try:
            connector = TCPConnector(limit=MAX_CONCURRENT, force_close=True, enable_cleanup_closed=True, ssl=ssl_context, family=socket.AF_INET)
            timeout = ClientTimeout(total=TOTAL_TIMEOUT)
            async with ClientSession(connector=connector, timeout=timeout, headers={"User-Agent": USER_AGENT}) as session:
                tasks = [check_website(session, domain, csv_manager) for domain in domains]
                await asyncio.gather(*tasks)
            logging.info(f"Completed chunk {chunk_number}")
            return
        except Exception as e:
            logging.error(f"Error processing chunk {chunk_number}: {str(e)}")
            if retry < CHUNK_RETRIES - 1:
                logging.info(f"Retrying chunk {chunk_number} in {CHUNK_RETRY_DELAY} seconds...")
                await asyncio.sleep(CHUNK_RETRY_DELAY)
            else:
                logging.error(f"Failed to process chunk {chunk_number} after {CHUNK_RETRIES} attempts")

async def process_domains(domains, processed_domains_set):
    global processed_domains
    csv_manager = CSVManager(OUTPUT_FILE, mode='w')  # Changed to 'w' mode

    domains_to_process = [domain for domain in domains if domain not in processed_domains_set]
    total_chunks = len(domains_to_process) // CHUNK_SIZE + (1 if len(domains_to_process) % CHUNK_SIZE > 0 else 0)

    for i in range(0, len(domains_to_process), CHUNK_SIZE):
        chunk = domains_to_process[i:i+CHUNK_SIZE]
        chunk_number = i//CHUNK_SIZE + 1
        logging.info(f"Processing chunk {chunk_number} of {total_chunks}")
        await process_chunk(chunk, csv_manager, chunk_number)
        logging.info(f"Pausing for 5 seconds before next chunk")
        await asyncio.sleep(5)  # Add a small pause between chunks

    csv_manager.close()

def read_processed_domains():
    processed_domains = set()
    domains_to_recheck = set()
    if os.path.exists(OUTPUT_FILE):
        with open(OUTPUT_FILE, 'r', newline='', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile)
            next(reader)  # Skip header
            for row in reader:
                if len(row) >= 2:
                    domain, status = row[0], row[1]
                    if status in ['yes', 'no']:
                        processed_domains.add(domain)
                    elif status == 'checking':
                        domains_to_recheck.add(domain)
    return processed_domains, domains_to_recheck

async def main():
    global total_domains, processed_domains
    logging.info("Starting alive-check.py")

    logging.info(f"Reading domains from {INPUT_FILE}")
    with open(INPUT_FILE, 'r', encoding='utf-8') as infile:
        reader = csv.reader(infile)
        all_domains = [row[0] for row in reader if row]

    processed_domains_set, domains_to_recheck = read_processed_domains()
    domains_to_process = set(all_domains) - processed_domains_set
    domains_to_process.update(domains_to_recheck)

    total_domains = len(all_domains)
    processed_domains = len(processed_domains_set)
    logging.info(f"Loaded {total_domains} domains")
    logging.info(f"Already processed: {processed_domains}")
    logging.info(f"Domains to process: {len(domains_to_process)}")

    start_time = time.time()
    await process_domains(list(domains_to_process), processed_domains_set)
    elapsed_time = time.time() - start_time

    with open(OUTPUT_FILE, 'r', encoding='utf-8') as outfile:
        reader = csv.reader(outfile)
        next(reader)  # Skip header
        results = list(reader)

    alive_count = sum(1 for row in results if row[1] == "yes")
    speed = len(domains_to_process) / elapsed_time if elapsed_time > 0 else 0
    logging.info(f"Processed {len(domains_to_process)} domains in {elapsed_time:.2f} seconds")
    logging.info(f"Found {alive_count} alive websites")
    logging.info(f"Average speed: {speed:.2f} domains/second")

    logging.info("alive-check.py completed")

if __name__ == "__main__":
    asyncio.run(main())