#! /usr/bin/python3

# https://github.com/dotnet/runtime/blob/v8.0.2/src/installer/managed/Microsoft.NET.HostModel/AppHost/HostWriter.cs#L170
# https://github.com/dotnet/runtime/blob/v8.0.2/src/installer/managed/Microsoft.NET.HostModel/Bundle/FileEntry.cs
# https://github.com/dotnet/runtime/blob/v8.0.2/src/installer/managed/Microsoft.NET.HostModel/Bundle/Manifest.cs

import argparse
import gettext
import json
import mmap
import os
import struct
import subprocess
import sys
from datetime import datetime

from packaging import version

translation = gettext.translation(
    "bsign-aliens", localedir="/usr/share/locale", fallback=True
)
_ = translation.gettext

bsign_version = subprocess.check_output(
    "bsign --version", shell=True, encoding="utf-8"
)[7:12]

bsign_codes = {
    0: _("no error"),
    1: _("permission denied"),
    6: _("failed to parse arguments"),
    7: _("failed to create temp file"),
    8: _("failed to write to temp file"),
    9: _("failed to parse ELF file"),
    10: _("failed to work with file xattr"),
    12: _("no memory"),
    21: _("is directory"),
    22: _("invalid argument"),
    24: _("too many open files"),
    26: _("file busy"),
    28: _("no space on device"),
    36: _("name too long"),
    64: _("no hash found"),
    65: _("no signature found"),
    66: _("bad hash found"),
    67: _("bad signature found"),
    68: _("unsupported file type"),
    69: _("bad pass phrase"),
    70: _("rewrite failed"),
    71: _("premature application termination"),
    72: _("exec failed because program wasn't found"),
    73: _("unused bytes after sig are not zero"),
}

FILE_OFFSET = 8
FILE_LENGTH = 8
FILE_NAME_LENGTH = 2
FILE_NAME_LENGTH_V7 = 1
DIGSIG_ELF_SIG_SIZE = 512

Elf64_Word = 4
Elf64_Xword = 8
Elf64_Addr = 8
Elf64_Off = 8

Elf32_Word = 4
Elf32_Addr = 4
Elf32_Off = 4


def process_message(mode, message):
    """ "
    Logging implementation for bsign-integrator.
    mode 0 - standard message to stdout
    mode 1 - message to /var/log/bsign-integrator.log
    """
    if mode == 0:
        print(message)
    else:
        current_time = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
        os.system(
            "echo \"{0} component=bsign-dot-net message='{1}'\" >> /var/log/bsign-integrator.log".format(
                current_time, message
            )
        )


def get_list_header_entry(mm, first_entry_offset):
    """
    Get entry header (file offset|file length|file name length|file name) list
    :param mm: memory map of opened file
    :param first_entry_offset: first entry offset in .net payload header
    :return: header entry list
    """
    list_entry_header = []
    mm.seek(first_entry_offset)
    header = mm.read()

    start_offset = 0
    additional_offset = 9
    length_offset = FILE_OFFSET + FILE_LENGTH
    name_offset = length_offset + additional_offset + FILE_NAME_LENGTH_V7
    while start_offset < len(header):
        file_off_and_length = struct.unpack(
            "<QQ", header[start_offset : start_offset + length_offset]
        )
        file_name_length = struct.unpack(
            ">B",
            header[
                start_offset
                + length_offset
                + additional_offset : start_offset
                + name_offset
            ],
        )
        file_name = header[
            start_offset
            + name_offset : start_offset
            + name_offset
            + file_name_length[0]
        ]
        list_entry_header.append(
            {
                "file-offset": file_off_and_length[0],
                "file-length": file_off_and_length[1],
                "file-name": file_name,
            }
        )
        start_offset += name_offset + len(list_entry_header[-1]["file-name"])

    return list_entry_header


def delete_files(*files):
    """
    Delete files
    :param files: files for delete
    """
    for file in files:
        if os.path.exists(file):
            os.remove(file)


def get_bundle_header_offset(file_name):
    try:
        with open(file_name, "rb") as file:
            mm = mmap.mmap(file.fileno(), 0, prot=mmap.PROT_READ)
            string = b"\x8b\x12\x02\xb9\x6a\x61\x20\x38\x72\x7b\x93\x02\x14\xd7\xa0\x32\x13\xf5\xb9\xe6\xef\xae\x33\x18\xee\x3b\x2d\xce\x24\xb3\x6a\xae"
            bundle_marker_offset = mm.read().rfind(string)
            mm.seek(bundle_marker_offset)
            mm.seek(-8, 1)
            bytes_data = mm.read(8)
            bundle_header_offset = int.from_bytes(bytes_data, byteorder="little")
            mm.close()
            return bundle_header_offset, bundle_marker_offset - 8
    except FileNotFoundError:
        process_message(mode, _("File not found!"))


def get_json_offsets(mm, opened_file):
    bundle_header_offset, bundle_header_marker = get_bundle_header_offset(opened_file)
    payload_list = get_list_header_entry(mm, bundle_header_offset + 85)
    for entry in payload_list:
        if b".deps.json" in entry["file-name"]:
            deps_offset = entry["file-offset"]
        if b".runtimeconfig.json" in entry["file-name"]:
            runtimeconfig_offset = entry["file-offset"]
    deps_bytes = struct.pack("<I", int(deps_offset))
    runtimeconfig_bytes = struct.pack("<I", int(runtimeconfig_offset))
    with open(opened_file, "rb") as file:
        data = file.read()
        deps_adress_offset = data.find(deps_bytes, bundle_header_offset)
        runtimeconfig_adress_offset = data.find(
            runtimeconfig_bytes, bundle_header_offset
        )
    return (
        deps_adress_offset,
        runtimeconfig_adress_offset,
        deps_offset,
        runtimeconfig_offset,
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "file", metavar="filename", help=_("path to elf file with .net payload to sign")
    )
    parser.add_argument(
        "-p", "--pgoptions", help=_("pass options to the privacy guard program")
    )
    parser.add_argument(
        "-c",
        "--check",
        help=_("check if file was successfully signed after script execution"),
        action="store_true",
    )
    parser.add_argument(
        "-i",
        "--integrator-mode",
        help=_("sign with integrator mode, debug/backend option"),
        action="store_true",
    )
    parser.add_argument(
        "-d",
        "--dotnet-check",
        help=_("check if ELF has .NET Bundle in it"),
        action="store_true",
    )
    parser.add_argument(
        "-v",
        "--dotnet-version",
        help=_("print .NET bundle version"),
        action="store_true",
    )
    args = parser.parse_args()
    opened_file = args.file
    new_file = opened_file + "_signed"

    if args.dotnet_version:
        result = subprocess.run(
            f'strings {opened_file} | grep "@(#)" | head -n 1 | sed -n \'s/^@(#)Version \(.*\) '
            f"@Commit:.*/\\1/p'",
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
        if result.stdout[:-1]:
            print(result.stdout[:-1])
            exit(0)
        else:
            isdotnet = subprocess.run(
                f"bsign-dot-net --dotnet-check {opened_file}",
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
            )
            if isdotnet.returncode == 0:
                print(_("Can't find dotnet version"))
                exit(1)
            else:
                print(isdotnet.stdout[:-1])
                exit(1)

    # Инструмент bsign, который мы будем использовать
    bsign = "bsign"

    # Переменная mode служит для определения режима работы - пользовательский - 0, bsign-integrator мод - 1
    global mode
    mode = 0
    if args.integrator_mode:
        # Для ведения логов необходим режим суперпользователя
        if not os.getuid() == 0:
            print(_("Must be sudo user to use in --integrator-mode"))
            exit(1)
        # Получаем значение инструмента bsign из конфигурации bsign-integrator
        try:
            with open("/etc/bsign-integrator/bsign-integrator.conf", "r") as file:
                config = json.load(file)
                bsign = config["default"]
        except:
            print(
                _(
                    "There is no /etc/bsign-integrator/bsign-integrator.conf of something wrong with it!"
                )
            )
            exit(1)
        mode = 1

    if version.parse(bsign_version) < version.parse("1.3.1"):
        process_message(mode, _("bsign version 1.3.1+ required!"))
        exit(1)

    # Проверяем, подписан ли файл. Если да, то подписывать его не нужно, это его поломает
    issigned = subprocess.run(
        f"bsign -wE {opened_file}",
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    if issigned.returncode == 0:
        process_message(mode, _("file already has signature in elf section!"))
        exit(1)

    try:
        with open(opened_file, "rb") as file:
            try:
                mm = mmap.mmap(file.fileno(), 0, prot=mmap.PROT_READ)
            except TypeError:
                process_message(mode, _("Mmap error!"))
                exit(1)
            else:
                try:
                    (
                        bundle_header_offset,
                        bundle_header_marker,
                    ) = get_bundle_header_offset(opened_file)
                except:
                    process_message(mode, _("File probably isn't ELF!"))
                    exit(1)
                if args.dotnet_check:
                    if bundle_header_offset:
                        process_message(mode, _("File has .NET bundle in it!"))
                        exit(0)
                    else:
                        process_message(mode, _("File hasn't .NET bundle in it!"))
                        exit(2)
                try:
                    payload_list = get_list_header_entry(mm, bundle_header_offset + 85)
                    (
                        deps_offset,
                        runtime_offset,
                        deps_adress_offset,
                        runtimeconfig_adress_offset,
                    ) = get_json_offsets(mm, opened_file)
                except struct.error:
                    process_message(
                        mode, _("File is already signed or wrong .NET bundle version!")
                    )
                    exit(1)
    except FileNotFoundError and IsADirectoryError:
        process_message(mode, _("No such file or directory: {0}").format(opened_file))
        exit(1)
    try:
        with open(opened_file, "rb") as original_file, open(
            new_file, "wb"
        ) as signed_file:
            # Перезаписываем Bundle Header Offset
            byte = original_file.read(bundle_header_marker)
            signed_file.write(byte)
            signed_file.write(struct.pack("<Q", bundle_header_offset + 592))
            original_file.seek(8, 1)
            # Перезаписываем JSON Deps
            byte = original_file.read(deps_offset - original_file.tell())
            signed_file.write(byte)
            signed_file.write(struct.pack("<I", int(deps_adress_offset) + 592))
            original_file.seek(len(struct.pack("<I", int(deps_adress_offset))), 1)

            # # Перезаписываем JSON Runtime
            byte = original_file.read(runtime_offset - original_file.tell())
            signed_file.write(byte)
            signed_file.write(struct.pack("<I", int(runtimeconfig_adress_offset) + 592))
            original_file.seek(
                len(struct.pack("<I", int(runtimeconfig_adress_offset) + 592)), 1
            )

            # Перезаписываем Bundle Manifest
            byte = original_file.read(
                (bundle_header_offset + 85) - original_file.tell()
            )
            signed_file.write(byte)

            for entry in payload_list:
                offset = original_file.read(8)
                length = original_file.read(8)
                decoded_offset = struct.unpack("<Q", offset)
                signed_file.write(struct.pack("<Q", decoded_offset[0] + 592))
                signed_file.write(length)

                other_stuff = original_file.read(10 + len(entry["file-name"]))
                signed_file.write(other_stuff)
    except PermissionError:
        process_message(mode, _("No permission to rewrite _signed file!"))
        sys.exit(1)

    bsign_args = [bsign, "--sign", "-E", new_file]
    if args.pgoptions is not None:
        if "--batch" in args.pgoptions:
            bsign_args.append("--nopass")
        bsign_args.extend(["-p", args.pgoptions])
    if subprocess.call(bsign_args):
        delete_files(new_file)
        process_message(mode, _("Error calling bsign failure"))
        return False

    name_pdp = subprocess.getstatusoutput("pdpl-file {0}".format(opened_file))
    if not name_pdp[0]:
        pdp_res = subprocess.getstatusoutput(
            "pdpl-file {0} {1}".format(name_pdp[1], new_file)
        )[0]
        if not pdp_res:
            process_message(
                mode, _("Created {0} with {1}").format(new_file, name_pdp[1])
            )
    else:
        process_message(
            mode, _("Cant put PDP on file {0}, created with 0:0").format(new_file)
        )
        process_message(
            mode, _("If you need to keep PDP label, rerun script with sudo")
        )

    if args.check:
        check_result = subprocess.run(
            ["bsign", "-cE", new_file], stdout=subprocess.DEVNULL
        )
        exit_code = check_result.returncode
        if exit_code == 0:
            print(_("Signature successfully found in {0}").format(new_file))
        elif exit_code < 64:
            print(
                _("Something went wrong while processing file {0}, {1}").format(
                    new_file, bsign_codes[exit_code]
                )
            )
        else:
            print(
                _("Signing file {0} failed, {1}").format(
                    new_file, bsign_codes[exit_code]
                )
            )
            sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    main()
