#!/usr/bin/env python3

# This code is in Public Domain. https://creativecommons.org/publicdomain/zero/1.0/legalcode
# Author: Nicholas Guriev <guriev-ns@ya.ru> 2021-12-26
# SPDX-License-Identifier:  CC0-1.0

"""
A tiny script to collect split debug data on MIPS. It can be reused on other
32-bit architectures if the need arises.

This script is a temporary replacement of dwp(1), DWARF packaging tool, because
the later one is crashing, see: <https://bugs.debian.org/1002601>. This script
is intended for use on a binary built with the `-gsplit-dwarf` compiler switch
and it has no safety checks if a debug data other than .dwo references is
present. You need at least GDB 11 or apply a one-liner patch to debug Telegram
after processing by dbgpatch.py. (GDB bug: <https://bugs.debian.org/1002602>)

See also: readelf(1), objcopy(1)
"""

import collections
import dataclasses
import math
import os
import shlex
import subprocess
import tempfile


@dataclasses.dataclass
class IndirectString:
    """References an indirect string as parsed from readelf(1) output."""

    offset: int
    value: str


@dataclasses.dataclass
class CompilationUnit:
    """Compilation unit in DWARF info parsed by readelf(1).

    This data-class describes only compilation unit produced with the
    `-gsplit-dwarf` compiler switch. It has two strings, a debug object name
    ending with .dwo and an actual compilation directory.
    """

    dwo_name: IndirectString
    comp_dir: IndirectString

    def __init__(self, dwo_name: tuple[int, str], comp_dir: tuple[int, str]):
        """Initialize a new compilation unit.

        This initializer automatically converts tuples describing debug strings.
        """

        self.dwo_name = IndirectString(*dwo_name)
        self.comp_dir = IndirectString(*comp_dir)

    @property
    def fullpath(self):
        """Full path to the split debug object.

        It is a concatenation of a compilation directory and an object name.
        """

        return os.path.join(self.comp_dir.value, self.dwo_name.value)


def parsedwarf(binary: str):
    """Parse DWARF data by means of readelf(1).

    Only comp_dir and dwo_name (with possible GNU extension of DWARF version 4)
    abbrevs are considered. Others do not remain after `--strip-dwo`.
    """

    look = lambda s: (s.split(':')[-1].strip(), s[s.index('<') + 1 : s.index('>')])
    with cmd('readelf', '-wi', '-wN', '--', binary) as output:
        dwo_name, comp_dir = None, None
        for line in output:
            if 'DW_AT_dwo_name' in line or 'DW_AT_GNU_dwo_name' in line:
                (dwo_name, off_dwo) = look(line)
            if 'DW_AT_comp_dir' in line:
                (comp_dir, off_dir) = look(line)
            if dwo_name and comp_dir:
                yield CompilationUnit(
                    (int(off_dwo, 16), dwo_name), (int(off_dir, 16), comp_dir))
                dwo_name, comp_dir = None, None


class Processing:
    """Main class responsible for processing debug data.

    It offers two methods, install and repack, which are intended to be called
    consecutively.
    """

    def __init__(self, binary: str):
        """Start processing.

        DWARF data of the `binary` is being consumed as part of initialization
        step.
        """

        self.__binary = binary
        self.__units = list(parsedwarf(binary))

    def install(self, target: str):
        """Do installation step.

        The method interprets the parsed list of compilation units and links
        supplementary .dwo files to the `target` directory. The directory is
        flattened in the result. If the binary consists of different objects
        with the same name, to avoid conflicts, they are being renamed plusing
        their ordinal number. For example, two `integrations.dwo` in different
        directories will become `integrations.1.dwo` and `integrations.2.dwo`.
        """

        mangled = collections.defaultdict(int)
        bases = collections.Counter(os.path.basename(u.fullpath)
                                    for u in self.__units)
        for unit in self.__units:
            base = os.path.basename(unit.fullpath)
            if bases[base] > 1:
                mangled[base] += 1
                (root, ext) = os.path.splitext(base)
                base = root + '.' + str(mangled[base]) + ext
            cmd.ln_f(unit.fullpath, os.path.join(target, base))
            unit.comp_dir.value, unit.dwo_name.value = target, base
        return self

    def repack(self, prefix: str):
        """Do repacking of the initial binary.

        After we moved debug objects to the installation directory, we should
        edit references in the binary so they point to actual files. First, we
        dump the .debug_info section with objcopy(1), then patch it and emit new
        content of the .debug_str section. After all, we replace these sections
        in the binary so that dh_strip(1) can find them.
        """

        with tempfile.TemporaryDirectory() as tmpdir:
            inf_path = os.path.join(tmpdir, 'debug_info.data')
            str_path = os.path.join(tmpdir, 'debug_str.data')
            cmd('objcopy', '--dump-section', f'.debug_info={inf_path}',
                '--', self.__binary).do

            with open(str_path, 'wb') as f:
                reloc = {u.comp_dir.offset: 0 for u in self.__units}
                address = f.write(prefix.encode() + b'\0')
                cmd.trace('printf', "'%s\\0'", shlex.quote(prefix),
                          '>', shlex.quote(str_path))
                for unit in self.__units:
                    reloc[unit.dwo_name.offset] = address
                    address += f.write(unit.dwo_name.value.encode() + b'\0')
                    cmd.trace('printf', "'%s\\0'", shlex.quote(unit.dwo_name.value),
                              '>>', shlex.quote(str_path))
            with open(inf_path, 'r+b') as f:
                width = math.ceil(math.log(max(reloc), 16))
                for offset, content in sorted(reloc.items()):
                    f.seek(offset)
                    f.write(content.to_bytes(4, 'little'))
                    cmd.trace('# write binary word 0x%08X at offset 0x%0*X to %s'
                              % (content, width, offset, shlex.quote(inf_path)))

            cmd('objcopy',
                '--update-section', f'.debug_str={str_path}',
                '--update-section', f'.debug_info={inf_path}',
                '--', self.__binary).do
        return self


class cmd:
    """A convenient wrapper around `subprocess.Popen()`.

    It simplifies a bit one-liner commands. It provides a method `.do` which you
    should call as a property.

    All commands are run with C locale to eliminate unreliable localized output.

    Read standard output in context manager.
    >>> with cmd('echo', 'Hello kitty!') as out:
    ...     out.read()
    'Hello kitty!\\n'

    Immediate execute command if no output is expected.
    >>> cmd('cat', '/dev/null').do

    Raises an error on non-zero exit code.
    >>> cmd('touch', '/untouchable').do
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    subprocess.CalledProcessError: Command '('touch', '/untouchable')' returned non-zero exit status 1.

    May be verbose and print which commands are being executed.
    >>> cmd.verbose = True
    >>> cmd('true', 'love stories', 'never have', 'endings').do
    true 'love stories' 'never have' endings
    """

    verbose = False
    """Print commands and their arguments before running?"""

    def __init__(self, *args: tuple[str, ...]):
        """Prepare a command for running."""

        self.__args = args
        self.__proc = None

    def __enter__(self):
        """Enter to context manager.

        Spawns a child process and returns a pipe whereto stdout is redirected.
        """

        self.__execute(True)
        return self.__proc.stdout

    def __exit__(self, exc_type, exc_value, traceback):
        """Exit from context manager.

        In case of an error, closes stdout to kill the process. Then waiting for
        a child process exits or stops.
        """

        if exc_type:
            self.__proc.stdout.close()
        else:
            self.__finish()

    @property
    def do(self):
        """Run the command in foreground.

        To execute the instantiated command immediately, you just appeal to this
        property. It is an indication that you are not opening a context
        manager but are calling to the command in void context.

        No output redirection is done. Waiting for a child process exits or
        stops.
        """

        self.__execute(False)
        self.__finish()

    @classmethod
    def ln_f(cls, src: str, dst: str):
        """Force link file.

        Create a hard link overriding existing destination if any. This may be
        considered as a built-in replacement for `cmd('ln', '-f', src, dst).do`
        given that `dst` is a valid link name.
        """

        cls.trace('ln', '-f', '--', shlex.quote(src), shlex.quote(dst))
        try:
            os.remove(dst)
        except FileNotFoundError:
            pass
        os.link(src, dst)

    @classmethod
    def trace(cls, *args: tuple[str, ...]):
        """Print args for tracing purposes.

        Or does nothing in silent mode.
        """

        if cls.verbose:
            print(*args)

    def __execute(self, readable: bool):
        self.trace(shlex.join(self.__args))
        extra = dict(stdout=subprocess.PIPE, text=True) if readable else {}
        self.__proc = subprocess.Popen(self.__args, env={'LANG': 'C'}, **extra)

    def __finish(self):
        self.__proc.wait()
        if rc := self.__proc.returncode:
            raise subprocess.CalledProcessError(rc, self.__args)


def main(argc: int, argv: list[str]):
    """Entry point to the present script."""

    if argc <= 1:
        print(f'usage: {argv[0]} [-x] <binary> <destination> <prefix>')
        return
    if argv[1] == '-x':
        cmd.verbose = True
        argv = argv[:1] + argv[2:]
    (binary, destination, prefix) = argv[1:]

    prefix = os.path.join('/', prefix)
    Processing(binary).install(destination + prefix).repack(prefix)


if __name__ == '__main__':
    import sys
    main(len(sys.argv), sys.argv)