query.py

#!/usr/bin/env python3

import argparse
import http
import json
import os
import requests
import sys
import time
import typing

from lib.detector.common import *

# Will be initialized by the ArgumentParser
ARGS: argparse.Namespace = None

with open("blacklist.json", "r", encoding="utf-8") as f:
    BLACKLIST = json.load(f)

with open("GITHUB_API_TOKEN", "r", encoding="utf-8") as f:
    GITHUB_API_TOKEN = f.readline().strip()

GITHUB_API_HEADERS = {
    "Accept": "application/vnd.github+json",
    "Authorization": f"Bearer {GITHUB_API_TOKEN}",
    "X-GitHub-Api-Version": "2022-11-28"
}

def api_get_json_generic(url: str, params: dict[str, typing.Any]) -> typing.Optional[dict]:
    loop = 60
    while True:
        try:
            response = requests.get(url,
                                    params = params,
                                    headers = GITHUB_API_HEADERS)
        except requests.exceptions.ConnectionError:
            print(f"[FAIL] Connection error. Sleeping for {loop} seconds...", file=sys.stderr)
            time.sleep(loop)
            loop *= 2
            continue

        match response.status_code:
            case 200: # OK
                return response.json()
            case 403: # Forbidden (= API rate limit reached)
                print(f"[FAIL] API rate limit reached. Sleeping for {loop} seconds...", file=sys.stderr)
                time.sleep(loop)
                loop *= 2
                continue
            case 404: # Not Found (= Project has been renamed)
                return None
            case 422: # Unprocessable Content (= Page Index greater than 20)
                print("[STOP] API page limit reached. Quitting...", file=sys.stderr)
                quit()
            case 500: # Internal Server Error (= Pull Request with no commits)
                return None

        # Non-reproducible and undocumented error codes
        print(f"[FAIL] {url} {params} returned code {response.status_code} ({http.client.responses[response.status_code]}). Sleeping for {loop} seconds...", file=sys.stderr)
        time.sleep(loop)
        loop *= 2
        continue

def api_search_repo(query: str, page: int) -> typing.Optional[dict]:
    return api_get_json_generic("https://api.github.com/search/repositories",
                                {"q": query,
                                 "sort": "stars",
                                 "per_page": ARGS.repos_per_page,
                                 "page": page + 1,
                                })

def api_repos_pulls(repository: str, page: int) -> typing.Optional[dict]:
    return api_get_json_generic(f"https://api.github.com/repos/{repository}/pulls",
                                {"state": "all",
                                 "sort": "created",
                                 "direction": "desc",
                                 "per_page": ARGS.pulls_per_page,
                                 "page": page + 1,
                                })

def api_pull_numbers(repository: str) -> list[int]:
    if ARGS.skip_pulls:
        return []
    if repository in BLACKLIST["skip_pulls"]:
        return []

    pull_numbers: list[int] = []

    for page in range(ARGS.pulls_page_offset, ARGS.pulls_page_offset + ARGS.pulls_pages):
        pulls_response = api_repos_pulls(repository, page)
        if (pulls_response is None) or (len(pulls_response) == 0):
            break

        pull_numbers.extend([pull["number"] for pull in pulls_response])

    return pull_numbers

def create_lock(repository: str) -> None:
    lock_path = os.path.join(ARGS.lock_directory, f"{repository}.lock")

    if not os.path.isdir(os.path.dirname(lock_path)):
        os.makedirs(os.path.dirname(lock_path))

    with open(lock_path, "w", encoding="utf-8") as f:
        print("lock", file=f)

def save_meta(repository, language, stars) -> None:
    with open(ARGS.qmeta_file, "a", encoding="utf-8") as f:
        print(f"{repository},{language},{stars}", file=f)

def init_argument_parser() -> None:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
    Copyright (C) 2025 Martin Weinzierl

    This program is free software; you can redistribute it and/or
    modify it under the terms of the GNU General Public License
    as published by the Free Software Foundation; version 2.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
""")

    parser.add_argument("--print-args", action="store_true", help="""
                         Print the given command line arguments to stderr
                        [default: disabled]""")

    parser.add_argument("--sanity-check", action="store_true", help="""
                         Prepend nickboucher/trojan-source to the list of repositories to check
                        [default: disabled]""")
    parser.add_argument("--repo-queries", type=str, nargs="+", default=["language:c"], help="""
                         The repository queries for the GitHub API
                        [default: language:c]""")

    parser.add_argument("--repos-per-page", type=int, default=10, help="""
                         The amount of items per repository page
                        [default: 10 items per page, max: 50]""")
    parser.add_argument("--pulls-per-page", type=int, default=10, help="""
                         The amount of items per pull request page
                        [default: 10 items per page, max: 50]""")

    parser.add_argument("--repo-pages", type=int, default=1, help="""
                         How many repositories should be checked per query
                        [default: 1 page]""")
    parser.add_argument("--pulls-pages", type=int, default=1, help="""
                         How many of the latest pull requests should be checked per repository
                        [default: 1 page]""")

    parser.add_argument("--repo-page-offset", type=int, default=0, help="""
                         How many repository pages should be skipped
                        [default: 0 pages]""")

    parser.add_argument("--pulls-page-offset", type=int, default=0, help="""
                         How many pull request pages should be skipped
                        [default: 0 pages]""")

    parser.add_argument("--skip-pulls", action="store_true", help="""
                         Skip pull request scanning
                        [default: disabled]""")

    parser.add_argument("--report-threshold", type=int, default=-1, help="""
                         Continue fetching pull requests if less than the given
                         number of reports exists for this repository (disabled: -1)
                        [default: -1]""")
    parser.add_argument("--stars-threshold", type=int, default=-1, help="""
                         Skip repositories with insufficient stars (disabled: -1)
                        [default: -1]""")

    parser.add_argument("--garbage-collect", action="store_true", help="""
                         Enable garbage collection in detect.py
                        [default: disabled]""")

    parser.add_argument("--lock-directory", type=str, default="locks", help="""
                         A directory path to store the lock files in
                        [default: locks]""")

    parser.add_argument("--qresult-file", type=str, default="qresults", help="""
                         A file to store the query results in
                        [default: qresults]""")
    parser.add_argument("--qmeta-file", type=str, default="qmeta.csv", help="""
                         A file to store the query metadata in
                        [default: qmeta.csv]""")


    global ARGS
    ARGS = parser.parse_args()

    if ARGS.print_args:
        print(ARGS, file=sys.stderr)

def main() -> None:
    init_argument_parser()

    repositories = set()
    qresult_map = {}

    # Fetch the most starred repositories matching the query given using the GitHub API
    for page in range(ARGS.repo_page_offset, ARGS.repo_page_offset + ARGS.repo_pages):
        for query in ARGS.repo_queries:
            print(f"Fetching repositories matching {query} (page {page})...", file=sys.stderr)
            repo_response = api_search_repo(query, page)

            if (repo_response is None) or (len(repo_response["items"]) == 0):
                print(f"[EMPTY] No repos. Query: {query}", file=sys.stderr)
                continue

            for item in repo_response["items"]:
                if (ARGS.stars_threshold != -1) and (item["stargazers_count"] < ARGS.stars_threshold):
                    print("Stars threshold satisfied. Skipping...", file=sys.stderr)
                    continue
                if (ARGS.report_threshold != -1) and (report_count(item["full_name"]) >= ARGS.report_threshold):
                    print("Report threshold satisfied. Skipping...", file=sys.stderr)
                    continue

                save_meta(item["full_name"], item["language"], item["stargazers_count"])
                repositories.add(item["full_name"])

    # Exit immediately if query returned no repositories
    if len(repositories) == 0:
        print("[EMPTY] No repos (all queries). Quitting.", file=sys.stderr)
        quit()

    # Add the demo repository to the list (sanity check)
    if ARGS.sanity_check:
        repositories.add("nickboucher/trojan-source")

    print(f"Found {len(repositories)} repositories", file=sys.stderr)

    # Exclude repositories mentioned in repository blacklist
    repositories = set(r for r in repositories if r not in BLACKLIST["skip_repo"])

    # Create locks and write default branch to qresult file
    with open(ARGS.qresult_file, "w", encoding="utf-8") as qresult_file:
        for repository in repositories:
            create_lock(repository)
            print(f"{repository} default", file=qresult_file)

    # Exclude repositories mentioned in pull request blacklist
    repositories = set(r for r in repositories if r not in BLACKLIST["skip_pulls"])

    # Store pull requests for every repository in a map
    for i, repository in enumerate(repositories):
        print(f"Fetching pull requests for repository {repository} ({i + 1}/{len(repositories)})...", file=sys.stderr)
        pull_numbers = api_pull_numbers(repository)

        # Make sure the key exists in the map
        qresult_map.setdefault(repository, [])

        if len(pull_numbers) == 0:
            print(f"[EMPTY] No pull. Repository: {repository}", file=sys.stderr)
            continue

        for pull_number in pull_numbers:
            qresult_map[repository].append(pull_number)

    # Append pull requests to qresult file
    with open(ARGS.qresult_file, "a", encoding="utf-8") as qresult_file:
        while True:
            # Stop if there are no more pull requests to append
            if len(qresult_map) == 0:
                break

            # Loop over the first few repositories
            for i in range(min(os.cpu_count(), len(qresult_map))):
                # Prevent index out of range
                if i >= len(repositories):
                    break

                repository = list(repositories)[i]

                # Restart outer loop
                if repository not in qresult_map.keys():
                    repositories.remove(repository)
                    break

                # Remove the entry if everything was popped
                if len(qresult_map[repository]) == 0:
                    qresult_map.pop(repository)
                    repositories.remove(repository)
                    break

                pull_number = qresult_map[repository].pop()

                print(f"{repository} pull {pull_number}", end="", file=qresult_file)

                # No more work to be done with this repository
                if ARGS.garbage_collect and (len(qresult_map[repository]) == 0):
                    print(" gc", file=qresult_file)
                    continue

                print("", file=qresult_file)

    assert len(qresult_map) == 0
    assert len(repositories) == 0

if __name__ == "__main__":
    main()