crawler.py

#!/usr/bin/env python3
"""
Recursively crawl a website and list all discovered URLs.

Also attempts to seed the crawl from sitemap(s) if found.

Usage:
    python crawler.py <url> [options]

Examples:
    python crawler.py https://example.com
    python crawler.py example.com --max-depth 5 --workers 8
    python crawler.py https://example.com --no-sitemap --delay 0.5
    python crawler.py https://example.com --output urls.txt
    python crawler.py https://example.com --include-subdomains
"""

import argparse
import re
import sys
import threading
import time
import xml.etree.ElementTree as ET
from collections import deque
from concurrent.futures import ThreadPoolExecutor, as_completed
from urllib.parse import urljoin, urlparse, urlunparse, urldefrag

import requests
from bs4 import BeautifulSoup

# Prefer lxml for speed and robustness if it's installed, else fall back.
try:
    import lxml  # noqa: F401
    BS_PARSER = "lxml"
except ImportError:
    BS_PARSER = "html.parser"

DOCTYPE_RE = re.compile(rb"<!DOCTYPE", re.IGNORECASE)

SITEMAP_NAMESPACES = {"sm": "http://www.sitemaps.org/schemas/sitemap/0.9"}

SITEMAP_PATHS = [
    "/sitemap.xml",
    "/sitemap_index.xml",
    "/sitemap-index.xml",
    "/wp-sitemap.xml",
    "/sitemap/sitemap.xml",
    "/sitemap1.xml",
]

DEFAULT_HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
        "(KHTML, like Gecko) Chrome/120.0 Safari/537.36"
    )
}

# Extensions we don't bother fetching HTML from
SKIP_EXTENSIONS = {
    ".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg", ".ico", ".bmp",
    ".css", ".js", ".mjs", ".json", ".xml", ".rss", ".atom",
    ".pdf", ".zip", ".gz", ".tar", ".rar", ".7z", ".bz2",
    ".mp3", ".mp4", ".avi", ".mov", ".mkv", ".webm", ".wav", ".ogg",
    ".woff", ".woff2", ".ttf", ".otf", ".eot",
    ".doc", ".docx", ".xls", ".xlsx", ".ppt", ".pptx",
    ".exe", ".dmg", ".pkg", ".deb", ".rpm", ".apk",
}

SCHEMES_TO_CRAWL = {"http", "https"}


# ---------- URL helpers ----------

def normalize_url(url: str) -> str:
    url = url.strip()
    if not url:
        raise ValueError("Empty URL provided.")
    if not urlparse(url).scheme:
        url = "https://" + url
    parsed = urlparse(url)
    # Strip fragment; lowercase scheme + host for consistency
    return urlunparse(parsed._replace(
        scheme=parsed.scheme.lower(),
        netloc=parsed.netloc.lower(),
        fragment="",
    ))


def canon(url: str) -> str:
    """Canonical form for dedup: drop fragment, normalize empty path -> '/'."""
    url, _ = urldefrag(url)
    p = urlparse(url)
    path = p.path or "/"
    return urlunparse(p._replace(path=path, params="", fragment=""))


def registrable_host(host: str) -> str:
    """Rough eTLD+1. Good enough for personal sites; not a full PSL impl."""
    host = host.lower().strip(".")
    parts = host.split(".")
    if len(parts) <= 2:
        return host
    # Handle common two-part TLDs
    two_part_tlds = {"co.uk", "com.au", "co.jp", "co.nz", "com.br", "co.in"}
    if ".".join(parts[-2:]) in two_part_tlds and len(parts) >= 3:
        return ".".join(parts[-3:])
    return ".".join(parts[-2:])


def host_allowed(url: str, base_host: str, include_subdomains: bool) -> bool:
    try:
        host = (urlparse(url).hostname or "").lower()
    except ValueError:
        return False
    base_host = base_host.lower()
    if include_subdomains:
        return registrable_host(host) == registrable_host(base_host)
    return host == base_host


def should_fetch_as_html(url: str) -> bool:
    """Filter out obviously non-HTML URLs by extension."""
    try:
        path = urlparse(url).path.lower()
    except ValueError:
        return False
    dot = path.rfind(".")
    slash = path.rfind("/")
    if dot > slash:  # extension in last segment
        ext = path[dot:]
        if ext in SKIP_EXTENSIONS:
            return False
    return True


# ---------- Sitemap parsing ----------

def _parse_sitemap_xml(content, session, timeout, depth=0, max_depth=3):
    urls = set()
    if depth > max_depth:
        return urls

    # Reject documents with DTDs outright. Sitemaps never legitimately
    # need one, and refusing them eliminates entity-expansion vectors.
    if DOCTYPE_RE.search(content):
        return urls

    try:
        root = ET.fromstring(content)
    except ET.ParseError:
        return urls

    tag = root.tag.split("}")[-1].lower() if "}" in root.tag else root.tag.lower()

    if tag == "sitemapindex":
        children = root.findall(".//sm:sitemap/sm:loc", SITEMAP_NAMESPACES) \
            or list(root.iter("loc"))
        for sitemap in children:
            if not sitemap.text:
                continue
            child_url = sitemap.text.strip()
            try:
                r = session.get(child_url, timeout=timeout)
                if r.status_code == 200:
                    urls |= _parse_sitemap_xml(
                        r.content, session, timeout, depth + 1, max_depth
                    )
            except requests.RequestException:
                continue
        return urls

    locs = root.findall(".//sm:loc", SITEMAP_NAMESPACES) or list(root.iter("loc"))
    for loc in locs:
        if loc.text:
            urls.add(loc.text.strip())
    return urls


def get_urls_from_sitemap(base_url, session, timeout=10):
    urls = set()
    for path in SITEMAP_PATHS:
        sitemap_url = urljoin(base_url, path)
        try:
            r = session.get(sitemap_url, timeout=timeout)
        except requests.RequestException:
            continue
        if r.status_code != 200:
            continue
        urls |= _parse_sitemap_xml(r.content, session, timeout)
        if urls:
            break
    return urls


# ---------- Crawler ----------

class Crawler:
    def __init__(
        self,
        base_url,
        session,
        timeout=10,
        delay=0.0,
        workers=8,
        max_depth=None,
        max_pages=None,
        include_subdomains=False,
        verbose=False,
    ):
        self.base_url = base_url
        self.base_host = (urlparse(base_url).hostname or "").lower()
        self.session = session
        self.timeout = timeout
        self.delay = delay
        self.workers = workers
        self.max_depth = max_depth
        self.max_pages = max_pages
        self.include_subdomains = include_subdomains
        self.verbose = verbose

        self.visited = set()        # canonical URLs we have fetched (or tried)
        self.discovered = set()     # all URLs we've seen (fetched + linked)
        self.lock = threading.Lock()
        self.rate_lock = threading.Lock()
        self.last_request = 0.0

    def _log(self, msg):
        if self.verbose:
            print(msg, file=sys.stderr)

    def _throttle(self):
        if self.delay <= 0:
            return
        with self.rate_lock:
            now = time.monotonic()
            wait = self.delay - (now - self.last_request)
            if wait > 0:
                time.sleep(wait)
            self.last_request = time.monotonic()

    def _fetch(self, url):
        self._throttle()
        try:
            r = self.session.get(url, timeout=self.timeout, allow_redirects=True)
        except requests.RequestException as e:
            self._log(f"[!] {url} -> {e}")
            return None
        # Record redirect target as discovered
        if r.url and r.url != url:
            final = canon(r.url)
            with self.lock:
                self.discovered.add(final)
        return r

    def _extract_links(self, html, page_url):
        soup = BeautifulSoup(html, BS_PARSER)
        out = set()
        for a in soup.find_all("a", href=True):
            href = a["href"].strip()
            if not href or href.startswith(("mailto:", "tel:", "javascript:", "data:")):
                continue
            full = urljoin(page_url, href)
            full = canon(full)
            p = urlparse(full)
            if p.scheme not in SCHEMES_TO_CRAWL:
                continue
            if not host_allowed(full, self.base_host, self.include_subdomains):
                continue
            out.add(full)
        return out

    def crawl(self, seeds):
        """BFS crawl. Returns the set of all discovered URLs."""
        # queue holds (url, depth)
        queue = deque()
        with self.lock:
            for s in seeds:
                s = canon(s)
                p = urlparse(s)
                if p.scheme not in SCHEMES_TO_CRAWL:
                    continue
                if not host_allowed(s, self.base_host, self.include_subdomains):
                    continue
                # Don't fetch known non-HTML seeds (images, PDFs, etc.)
                if not should_fetch_as_html(s):
                    continue
                if s in self.visited:
                    continue
                self.discovered.add(s)
                queue.append((s, 0))

        while queue:
            # Snapshot current frontier
            batch = []
            while queue:
                url, depth = queue.popleft()
                if url in self.visited:
                    continue
                if self.max_pages is not None and len(self.visited) >= self.max_pages:
                    queue.clear()
                    break
                self.visited.add(url)
                batch.append((url, depth))

            if not batch:
                break

            self._log(f"[*] Frontier: {len(batch)} URL(s) "
                      f"(visited so far: {len(self.visited)})")

            # Fetch batch concurrently
            results = []
            with ThreadPoolExecutor(max_workers=self.workers) as pool:
                future_map = {
                    pool.submit(self._fetch, url): (url, depth)
                    for url, depth in batch
                }
                for fut in as_completed(future_map):
                    url, depth = future_map[fut]
                    try:
                        r = fut.result()
                    except Exception as e:
                        self._log(f"[!] Worker error on {url}: {e}")
                        continue
                    results.append((url, depth, r))

            # Process responses
            new_items = []
            for url, depth, r in results:
                if r is None or r.status_code != 200:
                    continue
                ctype = r.headers.get("Content-Type", "").lower()
                if "html" not in ctype:
                    continue
                try:
                    links = self._extract_links(r.text, url)
                except Exception as e:
                    self._log(f"[!] Parse error on {url}: {e}")
                    continue

                with self.lock:
                    for link in links:
                        if link in self.discovered:
                            continue
                        self.discovered.add(link)
                        if not should_fetch_as_html(link):
                            continue
                        if self.max_depth is not None and depth + 1 > self.max_depth:
                            continue
                        new_items.append((link, depth + 1))

            for item in new_items:
                queue.append(item)

        return self.discovered


# ---------- CLI ----------

def build_arg_parser():
    p = argparse.ArgumentParser(
        description="Recursively crawl a website and list all discovered URLs.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("target", help="Target URL or hostname")
    p.add_argument("-t", "--timeout", type=float, default=10.0,
                   help="Per-request timeout in seconds (default: 10)")
    p.add_argument("-d", "--max-depth", type=int, default=None,
                   help="Max crawl depth from seeds (default: unlimited)")
    p.add_argument("-m", "--max-pages", type=int, default=None,
                   help="Max number of pages to fetch (default: unlimited)")
    p.add_argument("-w", "--workers", type=int, default=8,
                   help="Concurrent workers (default: 8)")
    p.add_argument("--delay", type=float, default=0.0,
                   help="Minimum delay between request starts, in seconds "
                        "(global, default: 0). Use ~0.2-1.0 to be polite.")
    p.add_argument("--no-sitemap", action="store_true",
                   help="Skip sitemap seeding and start from the given URL")
    p.add_argument("--include-subdomains", action="store_true",
                   help="Follow links to other subdomains of the same site")
    p.add_argument("-o", "--output", help="Write URLs to a file")
    p.add_argument("--insecure", action="store_true",
                   help="Disable TLS certificate verification")
    p.add_argument("-v", "--verbose", action="store_true",
                   help="Print progress to stderr")
    return p


def main():
    args = build_arg_parser().parse_args()

    try:
        site = normalize_url(args.target)
    except ValueError as e:
        print(f"[!] {e}", file=sys.stderr)
        return 2

    parsed = urlparse(site)
    if not parsed.netloc:
        print(f"[!] Could not parse host from: {args.target}", file=sys.stderr)
        return 2

    print(f"[*] Target: {site}", file=sys.stderr)

    session = requests.Session()
    session.headers.update(DEFAULT_HEADERS)
    session.verify = not args.insecure
    if args.insecure:
        import urllib3
        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        print("[!] TLS verification disabled.", file=sys.stderr)

    # Seed URLs
    seeds = {canon(site)}
    if not args.no_sitemap:
        print("[*] Checking for sitemap ...", file=sys.stderr)
        sitemap_urls = get_urls_from_sitemap(site, session, timeout=args.timeout)
        if sitemap_urls:
            print(f"[+] Sitemap yielded {len(sitemap_urls)} URL(s).",
                  file=sys.stderr)
            for u in sitemap_urls:
                seeds.add(canon(u))
        else:
            print("[*] No sitemap found.", file=sys.stderr)

    crawler = Crawler(
        base_url=site,
        session=session,
        timeout=args.timeout,
        delay=args.delay,
        workers=args.workers,
        max_depth=args.max_depth,
        max_pages=args.max_pages,
        include_subdomains=args.include_subdomains,
        verbose=args.verbose,
    )

    print(f"[*] Crawling with {args.workers} workers "
          f"(delay={args.delay}s) ...", file=sys.stderr)

    try:
        discovered = crawler.crawl(seeds)
    except KeyboardInterrupt:
        print("\n[!] Interrupted by user. Writing partial results.",
              file=sys.stderr)
        discovered = crawler.discovered

    if not discovered:
        print("[!] No URLs discovered.", file=sys.stderr)
        return 1

    output = "\n".join(sorted(discovered))

    if args.output:
        try:
            with open(args.output, "w", encoding="utf-8") as f:
                f.write(output + "\n")
            print(f"[+] Wrote {len(discovered)} URL(s) to {args.output}",
                  file=sys.stderr)
        except OSError as e:
            print(f"[!] Could not write output: {e}", file=sys.stderr)
            return 1
    else:
        print(output)

    print(f"[+] Fetched {len(crawler.visited)} page(s); "
          f"{len(discovered)} unique URL(s) discovered.", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())