analyse.py

#!/usr/bin/env python3

import argparse
import json
import multiprocessing
import os
import re
import subprocess
import sys
import time
import typing
import wrapt_timeout_decorator

from lib.detector.common import *

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

@wrapt_timeout_decorator.timeout(10)
def search_generic(data: str, pattern) -> typing.Optional[list[str]]:
    results = re.findall(pattern, data)

    if bool(results):
        # Filter empty matches and flatten list of tuples
        results = [item for tup in results for item in tup if len(item) > 0]
        if bool(results):
            return results

    return None

def search_invisible(data: str) -> typing.Optional[list[str]]:
    return search_generic(data, INVISIBLE_HEURISTICS_REGEX_PATTERN)

def search_homoglyph(data: str) -> typing.Optional[list[str]]:
    return search_generic(data, HOMOGLYPH_HEURISTICS_REGEX_PATTERN)

def launch_processes(files: list[str]) -> None:
    assert multiprocessing.get_start_method() == "fork"

    invisible_matches = 0
    homoglyph_matches = 0

    with multiprocessing.Pool(processes=ARGS.processes) as pool:
        result_tuples = list(pool.imap_unordered(process_main, files, chunksize=16))

    with open(ARGS.invisible_file, "w", encoding="utf-8") as invisible_file:
        with open(ARGS.homoglyph_file, "w", encoding="utf-8") as homoglyph_file:

            for invisible_out, homoglyph_out in result_tuples:
                if bool(invisible_out):
                    invisible_matches += 1
                    print(invisible_out, file=invisible_file)
                if bool(homoglyph_out):
                    homoglyph_matches += 1
                    print(homoglyph_out, file=homoglyph_file)

    print("----------", file=sys.stderr)
    print(f"Invisible matches: {invisible_matches} files", file=sys.stderr)
    print(f"Homoglyph matches: {homoglyph_matches} files", file=sys.stderr)

def process_main(file_path: str) -> tuple[typing.Optional[str], typing.Optional[str]]:

    if not os.path.isfile(file_path):
        # bool("") == False
        return ("", "")

    try:
        invisible_out, homoglyph_out = analyse_file(file_path)
    except TimeoutError:
        print(f"[TIMEOUT] {file_path}", file=sys.stderr)
        return ("", "")

    return (invisible_out, homoglyph_out)

def make_colorful(text: str) -> str:
    escaped_string = ""
    for char in text:
        codepoint = ord(char)

        # Printable ASCII character
        if 0x20 <= codepoint <= 0x7f:
            escaped_string += char
            continue

        show_codepoint = bool(INVISIBLE_REGEX_PATTERN.match(char))

        # Non-printable ASCII character or Unicode
        if codepoint <= 0xffff:
            if show_codepoint:
                escaped_string += f"{ANSI["fg_red"]}u{codepoint:04x}{ANSI["reset"]}"
            else:
                escaped_string += f"{ANSI["fg_red"]}{char}{ANSI["reset"]}"

            continue

        # Extended Unicode
        if show_codepoint:
            escaped_string += f"{ANSI["fg_red"]}U{codepoint:08x}{ANSI["reset"]}"
        else:
            escaped_string += f"{ANSI["fg_red"]}{char}{ANSI["reset"]}"

    return escaped_string

def list_matches(lines: list[str], matches: list[str]) -> str:
    duplicate_buffer = set()
    for i, line in enumerate(lines, start=1):
        for match in matches:
            if match in line:
                duplicate_buffer.add(f"[l. {i:04d}] : {line}")
                if ARGS.first_match_only:
                    break

    string_buffer = ""
    for e in duplicate_buffer:
        string_buffer += make_colorful(e) + "\n"

    return string_buffer

@wrapt_timeout_decorator.timeout(60)
def analyse_file(file_path: str) -> tuple[typing.Optional[str], typing.Optional[str]]:
    invisible_out = None
    homoglyph_out = None

    # Skip large files which are unlikely source code (> 1 MiB)
    if os.stat(file_path).st_size > 1_048_576:
        return (None, None)

    # Quick Fix: Skip VeraCrypt containers
    if os.path.splitext(file_path)[1] == ".hc":
        return (None, None)

    # The use of encoding "utf-8-sig" will skip an optional BOM
    # signature prepended to the actual file content
    encodings = ["utf-8-sig", "utf-16", "latin1", "iso-8859-1"]
    for encoding in encodings:
        with open(file_path, "r", encoding=encoding) as file:
            try:
                data = file.read()

                # Quick Fix: Extension *.ts may be a TypeScript file or an MPEG Transport Stream
                if os.path.splitext(file_path)[1] == ".ts":
                    if data.count("\x00") > 10:
                        return (None, None)

                # Remove comments from the source code,
                # in which Unicode codepoints are generally permitted.
                data = remove_comments(data, file_path)
                lines = data.splitlines()

                # Search for invisible codepoints
                try:
                    invisible_results = search_invisible(data)

                    if bool(invisible_results):
                        invisible_out = f"########## {file_path} ##########\n{list_matches(lines, invisible_results)}"

                except TimeoutError:
                    homoglyph_flag = False

                # Search for homoglyphs
                try:
                    homoglyph_results = search_homoglyph(data)

                    if bool(homoglyph_results):
                        homoglyph_out = f"########## {file_path} ##########\n{list_matches(lines, homoglyph_results)}"

                except TimeoutError:
                    homoglyph_flag = False

                # Explicitly free memory occupied by file contents
                del data
                del lines

                break

            except UnicodeDecodeError:
                # Print error message only, if every encoding failed
                if encoding == encodings[-1]:
                    print(f"{ANSI["bg_blu"]}{file_path}: Unicode Decode Error {encodings}{ANSI["reset"]}", file=sys.stderr)

    return (invisible_out, homoglyph_out)

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("--skip-pulls", action="store_true", help="""
                         Skip pull request scanning
                        [default: disabled]""")

    parser.add_argument("--first-match-only", action="store_true", help="""
                         Print the first match only
                        [default: disabled]""")

    parser.add_argument("--download-directory", type=str, default="downloads", help="""
                         A directory path to store the downloaded files in
                        [default: downloads]""")
    parser.add_argument("--report-directory", type=str, default="reports", help="""
                         A directory path to load the reports from
                        [default: reports]""")

    parser.add_argument("--bidi-file", type=str, default="bidirectional", help="""
                         A file path to store the previously BiDi-flagged files in
                        [default: bidirectional]""")
    parser.add_argument("--invisible-file", type=str, default="invisible", help="""
                         A file path to store the invisible-flagged files
                         and context for manual analysis in
                        [default: invisible]""")
    parser.add_argument("--homoglyph-file", type=str, default="homoglyph", help="""
                         A file path to store the homoglyph-flagged files
                         and context for manual analysis in
                        [default: homoglyph]""")

    parser.add_argument("--processes", type=int, default=os.cpu_count(), help=f"""
                         How many subprocesses to use
                        [default: os.cpu_count(), {os.cpu_count()} on this system]""")

    global ARGS
    ARGS = parser.parse_args()

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

def main() -> None:
    init_argument_parser()

    bidi_files = []
    unicode_files = []

    default_branch_file_counter = 0
    different_in_pull_counter = 0
    unchanged_in_pull_counter = 0
    total_files_counter = 0
    report_counter = 0

    print("Loading reports and comparing checksums...", file=sys.stderr)

    # Get all repositories in data set
    repositories = get_repositories(ARGS.report_directory)
    for i, repository in enumerate(repositories):

        default_branch = ""
        reports = {}

        # Load all reports and determine the default branch
        for report_file in get_json_files(os.path.join(ARGS.report_directory, repository)):
            report = load_report(report_file)
            report_counter += 1

            if report["is_default_branch"]:
                default_branch = report["branch"]

            reports[report["branch"]] = report

        if len(default_branch) == 0:
            print(f"[FAIL] {repository}: No default branch. Skipping repository...", file=sys.stderr)
            continue

        default_results = reports[default_branch]["results"]

        if ARGS.skip_pulls:
            reports = {default_branch: reports[default_branch]}

        # Loop over all reports of this repository
        for branch, report in reports.items():

            # Loop over all files in this report
            for absolute_file_path, flags in report["results"].items():
                total_files_counter += 1

                # Check if file was flagged in default branch too
                if (branch != default_branch) and (absolute_file_path in default_results):

                    # Compare checksums -> continue loop if they match
                    try:
                        if flags["sha256sum"] == default_results[absolute_file_path]["sha256sum"]:
                            unchanged_in_pull_counter += 1
                            continue
                    except KeyError:
                        pass

                if branch == default_branch:
                    default_branch_file_counter += 1
                else:
                    different_in_pull_counter += 1

                relative_file_path = absolute_file_path.split(f"/{report["repository"]}/")[-1]
                download_file_path = os.path.join(ARGS.download_directory, report["repository"], report["branch"], relative_file_path)

                if flags["bidi_flag"]:
                    bidi_files.append(download_file_path)
                if flags["unicode_flag"]:
                    unicode_files.append(download_file_path)

    print("----------", file=sys.stderr)
    print(f"Repositories: {len(repositories)}", file=sys.stderr)
    print(f"Reports: {report_counter}", file=sys.stderr)
    print("----------", file=sys.stderr)
    print(f"Suspicious: {total_files_counter} files", file=sys.stderr)
    print(f"Duplicates (by checksum): {unchanged_in_pull_counter} files", file=sys.stderr)
    print(f"Default branches: {default_branch_file_counter} files", file=sys.stderr)
    print(f"Different from default branch: {different_in_pull_counter} files", file=sys.stderr)
    print("----------", file=sys.stderr)
    print(f"BiDi matches: {len(bidi_files)} files (flagged by detect.py)", file=sys.stderr)
    print(f"Generic Unicode matches: {len(unicode_files)} files (to be analysed)", file=sys.stderr)
    print("----------", file=sys.stderr)

    print("Continue? (y/n) ", end="", file=sys.stderr)
    if input().lower() in ["y", "yes"]:
        print("This may take a while...", file=sys.stderr)

        with open(ARGS.bidi_file, "w", encoding="utf-8") as f:
            for file_path in bidi_files:
                print(file_path, file=f)

        launch_processes(unicode_files)

if __name__ == "__main__":
    main()