#! /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 os
import re
import shutil
import subprocess
import sys
from datetime import datetime

# from simple_manifest import SingleFileBundleParser
import dotnet_bundle as dnb
from bsign_helper import BSignHelper

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

# OFFSET_INCREMENT = 592
OFFSET_INCREMENT = 588


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 delete_files(*files):
    """
    Delete files
    :param files: files for delete
    """
    for file in files:
        if os.path.exists(file):
            os.remove(file)


def get_dotnet_version(opened_file):
    # First attempt: extract version from strings output
    result = subprocess.run(
        f"strings '{opened_file}' | grep '@(#)' | head -n 1 | sed -n 's/^@(#)Version \\(.*\\) @Commit:.*/\\1/p'",
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    if result.stdout and result.stdout.strip():
        version = result.stdout.strip()
        print(version)
        exit(0)

    # Second attempt: check if it's a dotnet file
    isdotnet = subprocess.run(
        ["bsign-dot-net", "--dotnet-check", opened_file],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    if isdotnet.returncode == 0:
        print(_("Can't find dotnet version"))
        exit(1)
    else:
        print(isdotnet.stdout.strip() if isdotnet.stdout else isdotnet.stderr.strip())
        exit(1)


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"

    # Инструмент 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.get("default")
                if not bsign:
                    print(_("Configuration file does not contain 'default' key"))
                    exit(1)
        except FileNotFoundError:
            print(
                _(
                    "There is no /etc/bsign-integrator/bsign-integrator.conf of something wrong with it!"
                )
            )
            exit(1)
        except json.JSONDecodeError as e:
            print(_("Invalid JSON in configuration file: {0}").format(str(e)))
            exit(1)
        mode = 1

    if not os.path.exists(opened_file):
        process_message(mode, _("No such file or directory: {0}").format(opened_file))
        exit(1)

    if args.dotnet_version:
        get_dotnet_version(opened_file)

    # Извлекаем пароль из pgoptions опций
    if args.pgoptions:
        passphrase_match = re.search(r'--passphrase=([^"\s]+)', args.pgoptions)
        passphrase = passphrase_match.group(1) if passphrase_match else None
    else:
        passphrase = None

    bsign = BSignHelper(passphrase=passphrase)

    # Проверяем доступность bsign
    available, msg = bsign.check_bsign_available("1.3.1")
    if not available:
        process_message(mode, msg)
        exit(1)

    bundle_parser = dnb.SingleFileBundleParser(opened_file)
    ret = bundle_parser.parse_bundle()

    if args.dotnet_check:
        if ret:
            process_message(mode, _("File has .NET bundle in it!"))
            exit(0)
        else:
            process_message(mode, _("File hasn't .NET bundle in it!"))
            exit(2)

    # Проверяем, подписан ли файл. Если да, то подписывать его не нужно, это его поломает
    if bsign.check_signature(opened_file):
        process_message(mode, _("file already has signature in elf section!"))
        exit(1)

    # Узнаем смещение, которое вызовет bsign
    signature_delta = bsign.calculate_signature_size(opened_file)

    try:
        shutil.copy2(opened_file, new_file)
    except PermissionError:
        print(f"Error: permission denied to write file '{new_file}'", file=sys.stderr)
        sys.exit(1)
    except FileNotFoundError:
        print(f"Error: source file not found '{opened_file}'", file=sys.stderr)
        sys.exit(1)
    except OSError as e:
        print(f"Error copying file: {e}", file=sys.stderr)
        sys.exit(1)

    # Поправим смещения (bsign увеличил elf заголовок):
    # 1. В заголовке бандла
    real_signature_offset = bundle_parser.header.signature_offset
    real_header_start = bundle_parser.header.manifest_header_start

    bundle_parser.header.signature_offset += signature_delta
    bundle_parser.header.manifest_header_start += signature_delta
    bundle_parser.header.manifest_header_end += signature_delta
    bundle_parser.header.deps_json_offset += signature_delta
    bundle_parser.header.runtime_config_offset += signature_delta
    dnb.BundleWriter.write_header(
        bundle_parser.header, new_file, real_signature_offset, real_header_start
    )

    # 2. В манифесте бандла
    real_section_start = bundle_parser.manifest.section_start

    bundle_parser.manifest.section_start += signature_delta
    bundle_parser.manifest.section_end += signature_delta
    for entry in bundle_parser.manifest.entries:
        entry.offset += signature_delta
    dnb.BundleWriter.write_manifest(
        bundle_parser.manifest,
        new_file,
        bundle_parser.header.major_version,
        real_section_start,
    )

    success, error_msg = bsign.sign_file(new_file, args.pgoptions)
    if not success:
        delete_files(new_file)
        process_message(mode, _(f"Error signing file: {error_msg}"))
        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:
        success, msg = bsign.verify_signature(new_file)
        if success:
            print(_("Signature successfully found in {0}").format(new_file))
        else:
            print(_("Signing file {0} failed, {1}").format(new_file, msg))
            sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    main()
