#! /usr/bin/python3

import argparse
import json
import multiprocessing
import os
import re
import subprocess
import sys

import magic

CONFFILE = "/etc/bsign-integrator/bsign-integrator.conf"
TOOLS_DICT = {"bsign": "/usr/bin/bsign"}

# коды выходов bsign с соответствующими им значениями, взяты из man bsign.
bsign_codes = {
    1: "доступ запрещен",
    2: "Некорректное использование инструмента",
    6: "не удалось спарсить аргументы",
    7: "не удалось создать временный файл",
    8: "не удалось записать во временный файл",
    9: "не удалось спарсить ELF-файл",
    10: "не удалось анализировать xattr файла",
    12: "не хватает оперативной памяти",
    21: "это папка",
    22: "неверный аргумент",
    24: "слишком много открытых файлов",
    26: "файл занят",
    28: "нет свободного места на устройстве",
    36: "слишком длинное имя",
    64: "хэш не найден",
    65: "подпись не найдена",
    66: "найден неверный хэш",
    67: "найдена неверная подпись",
    68: "неподдерживаемый формат файла",
    69: "неправильное ключевое слово",
    70: "не удалось перезаписать",
    71: "исполнение прервано",
    72: "исполнение не удалось, потому что программа не найдена",
    73: "неиспользуемые байты после подписи не равны нулю",
}


def get_cert_list_from_file(filename):
    try:
        with open(filename, "r") as file:
            lines = [line.strip() for line in file.readlines()]
            new_lines = []
            for line in lines:
                if line and (" " not in line):
                    if line not in parse_public_keys_id():
                        print(
                            f"Переданный сертификат - {line} - не импортирован, рассматриваться при анализе ФС не будет"
                        )
                    else:
                        new_lines.append(line)
            return new_lines
    except:
        print("Что-то пошло не так!")
        sys.exit(1)


def check_cert_in_file(file, tmpfile, additional_option):
    command = subprocess.run(
        f'bsign -V{additional_option} --pgoptions="--batch --pinentry-mode=loopback --no-default-keyring --keyring={tmpfile}" {file}',
        capture_output=True,
        text=True,
        shell=True,
    )
    if command.returncode == 0:
        return True
    else:
        return False


def analyse_file_list(file_list, tmpfilelist, additional_option):
    options = list(additional_option)
    good = False
    for file in file_list:
        for tmpfile in tmpfilelist:
            for option in options:
                good = False
                if check_cert_in_file(file, tmpfile, option):
                    print(file)
                    good = True
                    break
            if good:
                break


def parse_public_keys_id():
    gpg_keys = subprocess.run(
        ["gpg", "--list-keys"],
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        text=True,
    )
    key_ids = re.findall(r"\b[0-9A-F]{40}\b", gpg_keys.stdout)
    return key_ids


def execute_function_in_multiprocessing(num_processes, func, files_list, *args):
    processes = []

    size = len(files_list) // num_processes
    remainder = len(files_list) % num_processes

    parts = []
    start = 0
    for i in range(num_processes):
        part_size = size + 1 if i < remainder else size
        parts.append(files_list[start : start + part_size])
        start += part_size

    for n in range(0, num_processes):
        p = multiprocessing.Process(target=func, args=(parts[n], *args))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()


def get_file_list_from_file(filename):
    try:
        with open(filename, "r") as file:
            lines = [line.strip() for line in file.readlines()]
            new_lines = []
            for line in lines:
                if line and (" " not in line) and (is_file(line)):
                    new_lines.append(line)
            return new_lines
    except:
        print("Что-то пошло не так!")
        sys.exit(1)


def get_file_list_from_path(path):
    files = []
    if is_file(path):
        files.append(path)
        return files
    for dirpath, _, filenames in os.walk(path):
        for file in filenames:
            full_path = os.path.join(dirpath, file)
            files.append(full_path)
    return files


def sign_file_or_return_opts(filename, additional_options, additional_dotnet):
    return_options = []
    if is_elf(filename):
        # Если юзер не указывал опции вручную, то bsign-integrator сам решает, как подписывать
        if not (
            (
                "-A" in additional_options
                or "-D" in additional_options
                or "-X" in additional_options
                or "-E" in additional_options
            )
        ):
            if is_dot_net(filename):
                if (
                    subprocess.run(
                        f"bsign-dot-net --dotnet-version {filename}",
                        stdout=subprocess.PIPE,
                        text=True,
                        shell=True,
                    ).stdout[0]
                    == "8"
                ):
                    exit_code = subprocess.call(
                        f"bsign-dot-net {''.join(additional_dotnet)} {filename}",
                        shell=True,
                    )
                    if exit_code != 0:
                        print([exit_code])
                else:
                    version = subprocess.run(
                        f"bsign-dot-net --dotnet-version {filename}",
                        stdout=subprocess.PIPE,
                        text=True,
                        shell=True,
                    ).stdout
                    print(
                        f"Данная версия .NET - {version[:-1]} на данный момент, к сожалению, не поддерживается. Файл {filename} не подписан"
                    )
            else:
                return_options.append("-E")
    # Помимо проверки, является ли файл DLL, проверяем, не были ли указаны опции типа подписи
    elif is_dll(filename) and not (
        (
            "-A" in additional_options
            or "-D" in additional_options
            or "-X" in additional_options
            or "-E" in additional_options
        )
    ):
        return_options.append("-A")
    # По дефолту СКЗИ должен указывать подпись в расширенные атрибуты
    elif not (
        "-A" in additional_options
        or "-D" in additional_options
        or "-X" in additional_options
        or "-E" in additional_options
    ):
        return_options.append("-X")
    return return_options


def send_opts_to_sign_tool(filename, bsign_tool, bsign_option, additional_options):
    info = subprocess.run(
        f'{bsign_tool} {" ".join(bsign_option)} {" ".join(additional_options)} {filename}',
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        shell=True,
    )
    if info.stdout:
        print(info.stdout[:-1])
    if info.stderr:
        print(info.stderr[:-1])
    # Бывает, что bsign ничего не выводит, тогда нужно смотреть код выхода
    if not info.stderr and not info.stdout:
        print(bsign_codes[info.returncode], f"файл {filename} не подписан")
    return info.returncode


def sign_file_list(file_list, bsign_tool, additional_options, additional_dotnet):
    for file in file_list:
        opts = sign_file_or_return_opts(file, additional_options, additional_dotnet)
        if is_elf(file) and not (
            "-A" in additional_options
            or "-D" in additional_options
            or "-X" in additional_options
            or "-E" in additional_options
        ):
            if is_dot_net(file):
                continue
        exit_code = send_opts_to_sign_tool(
            file, bsign_tool, ["-s"], additional_options + opts
        )
        if exit_code >= 64:
            print(bsign_codes[exit_code], f"файл {file} не подписан")


def check_file_list(file_list, bsign_tool, additional_options):
    for file in file_list:
        info = subprocess.run(
            f'{bsign_tool} "-w" {" ".join(additional_options)} {file}',
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            shell=True,
        )
        print(info.stdout, info.stderr)


def custom_error_message(message):
    print(f"Ошибка: {message}")
    sys.exit(1)


def custom_help_message():
    print(
        "bsign_integrator [ДЕЙСТВИЕ] <ОПЦИИ> ...\n"
        "\n"
        "ДЕЙСТВИЯ:\n"
        "-b, --bsign                 Выводит значение криптографической утилиты, заданной по умолчанию)\n"
        "-d, --default [TOOL]        Задание криптографической утилиты TOOL, которая будет использоваться по умолчанию\n"
        "-w, --show-info [FILE]      Проверка подписи на заданном FILE\n"
        "-s, --sign [FILE]           Формирование подписи файла (по умолчанию, подпись записывается в FILE. Утилита самостоятельно определяет тип файла. Файлы типа ELF подписывает в спец секции, а остальные - в конец файла)\n"
        "-a, --sign-analyse [CERT]   Формирование списка файлов, подпись которых соответствует заданным параметрам. По умолчанию, выполняется поиск по файлам со всеми видами цифровой подписи.\n"
        "                            Файл попадает в список, если подпись файла соответствует SHA1 отпечатку сертификата ключа проверки ЦП, указанному в CERT.\n"
        "-T, --tools-list            Выводит список поддерживамых СКЗИ\n"
        "\n"
        "ОПЦИИ:\n"
        "-C, --current [PATH]    Задание криптографической утилиты по пути PATH, которая будет использована в текущей сессии интегратора\n"
        "-L, --input_list [FILE] В качестве входных значений использовать список, определённый в FILE\n"
        "-E, --elf-only          Подпись в elf-секции файла (только для файлов типа ELF)\n"
        "-A, --attached          Подпись в конце файла\n"
        "-X, --xattr-only        Подпись в расширенных атрибутах\n"
        "-D, --detached          Подпись в отдельном файле (отсоединённая подпись)\n"
        "-i, --include [PATH]    Добавить PATH к списку каталогов для поиска входных файлов. Опция может использоваться более одного раза\n"
        "-e, --exclude [PATH]    Удалить PATH из списка каталогов для поиска входных файлов. Опция может использоваться более одного раза\n"
        "-k, --key [KEY]         Задание ключа для использования в текущей сессии интегратора. Используется вместе с опцией -p/--passphrase\n"
        "-p, --passphrase [PASS] Задание пароля, соответствующего ключу, указанному в -k/--key\n"
        "-j, --jobs [NUMJOBS]    Количество процессов, которое разрешено запускать одновременно. Если не задано, то соответствует количеству процессоров"
    )


def get_config():
    with open(CONFFILE, "r") as file:
        config = json.load(file)
    return config


def is_elf(file_path):
    try:
        return "ELF" in magic.Magic().from_file(file_path)
    except:
        return False


def is_dll(file_path):
    return file_path.lower().endswith(".dll")


def is_file(file_path):
    if not os.path.isfile(file_path):
        return False
    return True


def is_dot_net(file_path):
    if (
        subprocess.call(
            f"bsign-dot-net --dotnet-check {file_path}",
            stdout=subprocess.DEVNULL,
            shell=True,
        )
        == 0
    ):
        return True
    else:
        return False


def exit_if_no_root():
    if os.geteuid() != 0:
        print(
            "Для использования данного функционала необходимо запустить инструмент с sudo"
        )
        sys.exit(1)


def set_default_tool(toolname):
    config = get_config()
    toolpath = TOOLS_DICT.get(toolname)
    if not toolpath:
        print(
            "Данный инструмент не поддерживается. Ознакомьтесь со списком доступных инструментов в bsign-integrator --tools-list"
        )
        sys.exit(1)
    config["bsign-default"] = toolpath
    with open(CONFFILE, "w") as file:
        json.dump(config, file, indent=4)
    sys.exit(0)


def main():
    additional_options = []
    bsign_option = []
    additional_dotnet = []
    paths = set()
    parser = argparse.ArgumentParser(add_help=False)

    parser.add_argument(
        "-b",
        "--bsign",
        action="store_true",
        help="Выводит значение криптографической утилиты, заданной по умолчанию",
    )
    parser.add_argument(
        "-d",
        "--default",
        nargs=1,
        type=str,
        metavar="PATH",
    )
    parser.add_argument(
        "-w",
        "--show-info",
        action="store_true",
    )
    parser.add_argument(
        "-s",
        "--sign",
        action="store_true",
    )
    parser.add_argument(
        "-a",
        "--sign-analyse",
        action="store_true",
    )

    parser.add_argument(
        "-C",
        "--current",
        nargs=1,
        metavar="PATH",
        type=str,
    )
    parser.add_argument(
        "-L",
        "--input-list",
        type=str,
        metavar="FILE",
        nargs=1,
    )
    parser.add_argument(
        "-E",
        "--elf-only",
        action="store_true",
    )
    parser.add_argument(
        "-A",
        "--attached",
        action="store_true",
    )
    parser.add_argument(
        "-X",
        "--xattr-only",
        action="store_true",
    )
    parser.add_argument(
        "-D",
        "--detached",
        action="store_true",
    )
    parser.add_argument(
        "FILE",
        type=str,
        nargs="*",
    )
    parser.add_argument(
        "-i",
        "--include",
        nargs="+",
        metavar="PATH",
        type=str,
        action="append",
    )
    parser.add_argument(
        "-e",
        "--exclude",
        nargs="+",
        metavar="PATH",
        type=str,
        action="append",
    )
    parser.add_argument(
        "-T",
        "--tools-list",
        action="store_true",
    )
    parser.add_argument(
        "-k",
        "--key",
        nargs=1,
        type=str,
        metavar="KEY",
    )
    parser.add_argument(
        "-p",
        "--passphrase",
        nargs=1,
        type=str,
        metavar="PASS",
    )
    parser.add_argument(
        "-j",
        "--jobs",
        nargs=1,
        type=int,
        metavar="NUMJOBS",
    )

    if "--help" in sys.argv[1:] or "-h" in sys.argv[1:]:
        custom_help_message()
        parser.exit()

    parser.error = custom_error_message

    args = parser.parse_args()
    try:
        filename = args.FILE[0]
    except IndexError:
        filename = None

    if args.sign:
        exit_if_no_root()
    if args.bsign:
        config = get_config()
        print(config["bsign-default"])
        sys.exit(0)
    if args.default:
        exit_if_no_root()
        set_default_tool(args.default[0])
    if (os.geteuid() == 0) and args.sign:
        if not args.passphrase:
            print(
                "Внимание! При использовании инструмента с sudo необходимо явно прописывать --passphrase и --key"
            )
            sys.exit(1)
    if args.tools_list:
        print("Список доступных инструментов:\n")
        version = subprocess.run(
            ["bsign --version | grep bsign"],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            shell=True,
        )
        if not version.stderr:
            print(
                "bsign, установлен, версия - {0}Стандартный инструмент для подписи".format(
                    version.stdout
                )
            )
        else:
            print("bsign, не установлен. Стандартный инструмент для подписи")
        sys.exit(0)
    if args.elf_only:
        additional_options.append("-E")
    elif args.attached:
        additional_options.append("-A")
    elif args.xattr_only:
        additional_options.append("-X")
    elif args.detached:
        if os.geteuid() != 0:
            print(
                "Внимание! Для успешного создания detached подписи необходимо запустить скрипт с sudo"
            )
        additional_options.append("-D")
    additional_options.append("-N")
    if (args.include or args.exclude) and os.geteuid() != 0:
        print("Опции --include и --exclude рекомендуется использовать с sudo")
    if args.include:
        for path in args.include:
            paths.update(get_file_list_from_path(path[0]))
    if args.exclude:
        for path in args.exclude:
            paths.difference_update(get_file_list_from_path(path[0]))
    if (args.include or args.exclude) and len(paths) == 0:
        print("Использованы опции -i/-e, но не были переданы аргументы")
        sys.exit(1)
    if (args.sign or args.show_info) and len(paths) == 0 and not args.input_list:
        if not filename:
            print("Не был передан FILE!")
            sys.exit(1)
        if not is_file(filename):
            print(f"{filename} не найден или не является файлом!")
            sys.exit(1)
    if args.key:
        if args.key[0]:
            if args.passphrase and args.passphrase[0]:
                additional_options.append(
                    f'--pgoptions="--batch --pinentry-mode=loopback --default-key={args.key[0]} --passphrase={args.passphrase[0]}"'
                )
                additional_dotnet.append(
                    f'--pgoptions="--batch --pinentry-mode=loopback --default-key={args.key[0]} --passphrase={args.passphrase[0]}"'
                )
            else:
                print("Не был передан --passphrase=PASS")
                sys.exit(1)
        else:
            print("--key: не был передан KEY")
            sys.exit(1)
    elif args.passphrase:
        if args.passphrase[0]:
            additional_options.append(
                f'--pgoptions="--batch --pinentry-mode=loopback --passphrase={args.passphrase[0]}"'
            )
            additional_dotnet.append(
                f'--pgoptions="--batch --pinentry-mode=loopback --passphrase={args.passphrase[0]}"'
            )
        else:
            print("--key: не был передан PASS")
            sys.exit(1)
    if args.show_info:
        bsign_option.append("-w")
    elif args.sign:
        # Здесь подписываем сингл-файл
        if not args.input_list and len(paths) == 0:
            opts = sign_file_or_return_opts(
                filename, additional_options, additional_dotnet
            )
            if (
                is_elf(filename)
                and is_dot_net(filename)
                and (
                    not (
                        (
                            "-A" in additional_options
                            or "-D" in additional_options
                            or "-X" in additional_options
                            or "-E" in additional_options
                        )
                    )
                )
            ):
                sys.exit(0)
            additional_options += opts
        bsign_option.append("-s")
    if args.jobs:
        num_processes = args.jobs[0]
    else:
        num_processes = os.cpu_count()
    if args.sign or args.show_info or args.sign_analyse:
        if not args.current:
            bsign_tool = get_config()["bsign-default"]
        else:
            if args.current[0] == "bsign":
                bsign_tool = args.current[0]
            else:
                print(
                    f"Инструмент {args.current[0]} не поддерживается. Ознакомьтесь со списком доступных инструментов в bsign-integrator --tools-list"
                )
                sys.exit(1)
    if args.show_info or args.sign:
        if args.input_list:
            filename = args.input_list[0]
            if not filename:
                print("FILE не был передан!")
                sys.exit(1)
            if not is_file(filename):
                print(f"{filename} не является файлом!")
                sys.exit(1)
            files_list = get_file_list_from_file(filename)
            if len(files_list) == 0:
                print("--input-list: Не удалось найти файлы!")
                sys.exit(1)
            if args.sign:
                execute_function_in_multiprocessing(
                    num_processes,
                    sign_file_list,
                    files_list,
                    bsign_tool,
                    additional_options,
                    additional_dotnet,
                )
            elif args.show_info:
                execute_function_in_multiprocessing(
                    num_processes,
                    check_file_list,
                    files_list,
                    bsign_tool,
                    additional_options,
                )
            sys.exit(0)
        if len(paths) > 0:
            if bsign_option[0] == "-s":
                execute_function_in_multiprocessing(
                    num_processes,
                    sign_file_list,
                    list(paths),
                    bsign_tool,
                    additional_options,
                    additional_dotnet,
                )
            elif bsign_option[0] == "-w":
                execute_function_in_multiprocessing(
                    num_processes,
                    check_file_list,
                    list(paths),
                    bsign_tool,
                    additional_options,
                )
            sys.exit(0)
        exit_code = send_opts_to_sign_tool(
            filename, bsign_tool, bsign_option, additional_options
        )
        sys.exit(exit_code)
    if args.sign_analyse:
        exit_if_no_root()
        tmpfilelist = []
        certlist = []
        if args.input_list:
            certlist = get_cert_list_from_file(args.input_list[0])
        else:
            certlist.append(filename)
        for cert in certlist:
            if not args.input_list and cert not in parse_public_keys_id():
                print(f"Переданный сертификат - {cert} - не импортирован")
                sys.exit(1)
            command = subprocess.run(
                ["mktemp", "/tmp/tempfile.XXXXXX"], capture_output=True, text=True
            )
            tmpfile = command.stdout.strip()
            subprocess.run(f"gpg --export {cert} > {tmpfile}", shell=True)
            tmpfilelist.append(tmpfile)

        additional_option = ""
        fs_filelist = get_file_list_from_path("/")

        if args.elf_only:
            additional_option += "E"
        if args.attached:
            additional_option += "A"
        if args.xattr_only:
            additional_option += "X"
        if args.detached:
            additional_option += "D"
        if additional_option == "":
            additional_option = "AEXD"

        execute_function_in_multiprocessing(
            num_processes,
            analyse_file_list,
            fs_filelist,
            tmpfilelist,
            additional_option,
        )
        for tmpfile in tmpfilelist:
            subprocess.run(f"rm {tmpfile}", shell=True)
        sys.exit(0)
    custom_help_message()
    sys.exit(1)


if __name__ == "__main__":
    main()
