main.py

import random
import os
import re
import sys
import time
import queue
import signal
import requests
from typing import List, Set
from bs4 import BeautifulSoup

BASE_URL = "https://github.com"
HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
    "Accept": "text/html,application/xhtml+xml",
    "Accept-Language": "en-US,en;q=0.9",
}

ROOT_DIR = os.path.join(os.getcwd(), "avatars")
USERNAME_RE = re.compile(r"^[a-zA-Z0-9-]{1,39}$")
REQUEST_TIMEOUT = 10
REQUEST_RETRIES = 2
SLEEP_BETWEEN_REQUESTS = 1.0
MAX_VISITED_HARD_LIMIT = 5000

visited = set()  # global set to track visited users


# --------------------------
# Signal handler
# --------------------------
def _signal_handler(sig, frame):
    print("\nInterrupted by user. Exiting...")
    sys.exit(0)


signal.signal(signal.SIGINT, _signal_handler)


# --------------------------
# Terminal colors
# --------------------------
try:
    import colorama

    colorama.init()
    RESET = colorama.Style.RESET_ALL
    BRIGHT = colorama.Style.BRIGHT
    RED = colorama.Fore.RED + BRIGHT
    GREEN = colorama.Fore.GREEN + BRIGHT
    YELLOW = colorama.Fore.YELLOW + BRIGHT
    CYAN = colorama.Fore.CYAN + BRIGHT
except Exception:
    RESET = "\033[0m"
    BRIGHT = "\033[1m"
    RED = "\033[1;31m"
    GREEN = "\033[1;32m"
    YELLOW = "\033[1;33m"
    CYAN = "\033[1;36m"


# --------------------------
# Prints
# --------------------------
def header():
    art = r"""
  _____ _ _   _       _     _            _             
 / ____(_) | (_)     | |   | |          | |            
| |  __ _| |_ _  __ _| |__ | | ___   ___| | _____ _ __ 
| | |_ | | __| |/ _` | '_ \| |/ _ \ / __| |/ / _ \ '__|
| |__| | | |_| | (_| | | | | | (_) | (__|   <  __/ |   
 \_____|_|\__|_|\__, |_| |_|_|\___/ \___|_|\_\___|_|   
                  __/ |                                 
                 |___/     
                 
                                              """
    print(f"{CYAN}{art}{RESET}")
    print(f"{BRIGHT}GitHub Avatar Crawler{RESET}")


def info(msg: str):
    print(f"{CYAN}[i]{RESET} {msg}")


def success(msg: str):
    print(f"{GREEN}✔{RESET} {msg}")


def warn(msg: str):
    print(f"{YELLOW}!{RESET} {msg}")


def error(msg: str):
    print(f"{RED}✖{RESET} {msg}")

# --------------------------
# Network helpers
# --------------------------
def safe_get(url: str):
    for attempt in range(REQUEST_RETRIES + 1):
        try:
            r = requests.get(url, headers=HEADERS, timeout=REQUEST_TIMEOUT)
            return r
        except requests.RequestException as e:
            warn(f"Request error {url}: {e} (attempt {attempt + 1})")
            time.sleep(random.uniform(0.6, 1.8))

    return None


def download_avatar(username: str, folder: str, filename: str) -> bool:
    url = f"{BASE_URL}/{username}"
    r = safe_get(url)
    if not r or r.status_code != 200:
        warn(f"Failed to fetch profile page for {username}")
        return False

    soup = BeautifulSoup(r.text, "html.parser")
    meta = soup.find("meta", property="og:image")
    if not meta or not meta.get("content"):
        warn(f"No avatar found for {username}")
        return False

    img_url = meta["content"].split("?", 1)[0]
    r2 = safe_get(img_url)
    if not r2 or r2.status_code != 200:
        warn(f"Failed to download avatar for {username}")
        return False

    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, filename)
    with open(path, "wb") as f:
        f.write(r2.content)
    success(f"Saved {path}")
    return True


def get_users_from_page(username: str, mode: str, visited: Set[str]) -> List[str]:
    users = set()
    page = 1
    while True:
        url = f"{BASE_URL}/{username}?page={page}&tab={mode}"
        r = safe_get(url)
        if not r or r.status_code != 200:
            break
        soup = BeautifulSoup(r.text, "html.parser")
        links = soup.select('a[data-hovercard-type="user"]')
        if not links:
            break
        for a in links:
            href = a.get("href", "")
            if href.startswith("/"):
                u = href.strip("/").split("?", 1)[0]
                if u and "/" not in u and u not in visited:
                    users.add(u)
        page += 1
        time.sleep(SLEEP_BETWEEN_REQUESTS)
    return list(users)


# --------------------------
# Structured scraping
# --------------------------
def scrape_structured(target: str, mode_choices: List[str], max_depth: int):
    info("Starting structured scraping...")
    base = os.path.join(ROOT_DIR, target)
    os.makedirs(base, exist_ok=True)
    download_avatar(target, base, "avatar.jpg")
    visited.add(target)

    for mode_choice in mode_choices:
        mode_dir = os.path.join(base, mode_choice)
        os.makedirs(mode_dir, exist_ok=True)
        first_level = get_users_from_page(target, mode_choice, visited)
        for user in first_level:
            if max_depth == 1:
                download_avatar(user, mode_dir, f"{user}.jpg")
                visited.add(user)
                time.sleep(SLEEP_BETWEEN_REQUESTS)
            else:
                user_folder = os.path.join(mode_dir, user)
                download_avatar(user, user_folder, "avatar.jpg")
                visited.add(user)
                time.sleep(SLEEP_BETWEEN_REQUESTS)
                _scrape_structured_recursive(user, mode_choice, 2, max_depth, user_folder)


def _scrape_structured_recursive(username: str, mode_choice: str, depth: int, max_depth: int, user_parent_dir: str):
    if depth > max_depth or len(visited) >= MAX_VISITED_HARD_LIMIT:
        return

    children_dir = os.path.join(user_parent_dir, mode_choice)
    os.makedirs(children_dir, exist_ok=True)

    children = get_users_from_page(username, mode_choice, visited)
    for child in children:
        if child in visited:
            continue
        if depth == max_depth:
            download_avatar(child, children_dir, f"{child}.jpg")
            visited.add(child)
            time.sleep(SLEEP_BETWEEN_REQUESTS)
        else:
            child_folder = os.path.join(children_dir, child)
            download_avatar(child, child_folder, "avatar.jpg")
            visited.add(child)
            time.sleep(SLEEP_BETWEEN_REQUESTS)
            # Recurse into both followers and following if 'both' mode
            for m in ["followers", "following"]:
                _scrape_structured_recursive(child, m, depth + 1, max_depth, child_folder)


# --------------------------
# Flat scraping
# --------------------------
def scrape_flat(target: str, mode_choices: List[str], max_depth: int):
    info("Starting flat scraping...")
    flat_dir = os.path.join(ROOT_DIR, target, "flatlist")  # <- flatlist folder added
    os.makedirs(flat_dir, exist_ok=True)

    q = queue.Queue()
    q.put((target, 0))
    download_avatar(target, flat_dir, f"{target}.jpg")
    visited.add(target)
    time.sleep(SLEEP_BETWEEN_REQUESTS)

    while not q.empty():
        username, depth = q.get()
        if depth >= max_depth:
            continue
        for mode_choice in mode_choices:
            children = get_users_from_page(username, mode_choice, visited)
            for child in children:
                if child in visited:
                    continue
                download_avatar(child, flat_dir, f"{child}.jpg")  # save in flatlist folder
                visited.add(child)
                time.sleep(SLEEP_BETWEEN_REQUESTS)
                if depth + 1 < max_depth:
                    q.put((child, depth + 1))


# --------------------------
# User input helpers
# --------------------------
def prompt_token() -> str:
    token = input("Do you want to use a GitHub Bearer Token? (leave blank for none): ").strip()
    if token:
        info("Token will be used for authenticated requests.")
    return token

def prompt_username() -> str:
    while True:
        u = input("GitHub username: ").strip()
        if USERNAME_RE.match(u):
            return u
        warn("Invalid username format.")


def prompt_mode() -> List[str]:
    while True:
        print("Choose mode:")
        print("  1) followers")
        print("  2) following")
        print("  3) both")
        choice = input("Choice (1/2/3): ").strip()
        if choice == "1":
            return ["followers"]
        if choice == "2":
            return ["following"]
        if choice == "3":
            return ["followers", "following"]
        warn("Invalid choice.")


def prompt_depth() -> int:
    while True:
        try:
            d = int(input("Enter depth (0=only self): ").strip())
            if d >= 0:
                return d
        except ValueError:
            pass
        warn("Depth must be a non-negative integer.")


def prompt_structure_choice() -> str:
    while True:
        print("Output format:")
        print("  1) Structured folder tree")
        print("  2) Flat folder")
        choice = input("Choose 1/2: ").strip()
        if choice in ["1", "2"]:
            return choice
        warn("Invalid choice.")


# --------------------------
# Main
# --------------------------
def main():
    header()
    target = prompt_username()
    token = prompt_token()
    if token:
        HEADERS["Authorization"] = f"Bearer {token}"

    modes = prompt_mode()
    depth = prompt_depth()

    structure_choice = prompt_structure_choice()

    info(f"Target={target}, Mode={modes}, Depth={depth}")
    if structure_choice == "1":
        scrape_structured(target, modes, depth)
    else:
        scrape_flat(target, modes, depth)

    success("Scraping complete. Check avatars/ folder.")


if __name__ == "__main__":
    main()