#!/usr/bin/env python3
# encoding=UTF-8

# Copyright © 2010-2019 Jakub Wilk <jwilk@jwilk.net>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the “Software”), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import argparse
import io
import os
import re
import subprocess
import sys

...  # Python 3 is required

__version__ = '0.2'

description = \
'''
print ELF dependencies in Graphviz format
'''

os.putenv('LC_ALL', 'C')

template_begin = '''
digraph {
    rankdir=LR
    edge [arrowsize=0.5, arrowhead="vee"]
    node [fontsize=10, width=0, height=0, shape=box]
'''

template_end = '''
}
'''

class Subprocess(subprocess.Popen):

    def __init__(self, *args, **kwargs):
        command_line = kwargs.get('args')
        if command_line is None:
            command_line = args[0]
        self.__command = command_line[0]
        subprocess.Popen.__init__(self, *args, **kwargs)

    def wait(self, *args, **kwargs):
        rc = subprocess.Popen.wait(self, *args, **kwargs)
        if rc != 0:
            raise subprocess.CalledProcessError(rc, self.__command)

def _html_escape(match):
    c = ord(match.group(0))
    return '&#{0};'.format(c)

def dot_escape(s):
    # https://www.graphviz.org/doc/info/lang.html says:
    #
    #    In quoted strings […] the only escaped character is double-quote
    #    ("). That is, in quoted strings, the dyad \" is converted to "; all
    #    other characters are left unchanged. In particular, \\ remains \\.
    #
    # This makes little sense, and doesn't match how dot(1) actually works.
    # Oh well. Doubling the backslashes can't hurt much.
    s = s.replace('\\', '\\\\')
    s = re.sub('[&"]', _html_escape, s)
    return '"{s}"'.format(s=s)

_readelf_needed = re.compile(r'^ 0x[0-9a-f]+ [(]NEEDED[)] +Shared library: \[(.*)\]$').match

def process(elf, mapping):
    print('   ', dot_escape(elf), '-> {')
    readelf = Subprocess(
        ['readelf', '-d', '--', elf],
        stdout=subprocess.PIPE
    )
    for line in readelf.stdout:
        line = line.decode('ASCII')
        match = _readelf_needed(line)
        if match is None:
            continue
        target = match.group(1)
        try:
            target = mapping[target]
        except KeyError:
            pass
        print('       ', dot_escape(target))
    readelf.wait()
    print('    }')

class VersionAction(argparse.Action):
    '''
    argparse --version action
    '''

    def __init__(self, option_strings, dest=argparse.SUPPRESS):
        super().__init__(
            option_strings=option_strings,
            dest=dest,
            nargs=0,
            help='show version information and exit'
        )

    def __call__(self, parser, namespace, values, option_string=None):
        print('{prog} {0}'.format(__version__, prog=parser.prog))
        print('+ Python {0}.{1}.{2}'.format(*sys.version_info))
        for prog in 'ldd', 'readelf':
            proc = Subprocess(
                [prog, '--version'],
                stdout=subprocess.PIPE
            )
            for line in proc.stdout:
                line = line.rstrip()
                line = line.decode('ASCII')
                print('+ ' + line)
                break
            proc.stdout.read()
            proc.wait()
        parser.exit()

def main():
    ap = argparse.ArgumentParser(description=description)
    ap.add_argument('--version', action=VersionAction)
    ap.add_argument('binary', metavar='BINARY', help='ELF binary to inspect')
    options = ap.parse_args()
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='UTF-8')
    try:
        os.stat(options.binary)
    except OSError as exc:
        print('{prog}: {path}: {err}'.format(prog=ap.prog, path=options.binary, err=exc.strerror), file=sys.stderr)
        sys.exit(1)
    ldd = Subprocess(
        ['ldd', '--', options.binary],
        stdout=subprocess.PIPE,
    )
    mapping = {}
    for line in ldd.stdout:
        line = line.decode('ASCII')
        if not line.startswith('\t'):
            continue
        if not line.endswith('\n'):
            continue
        line = line[1:-1].split(' => ')
        if len(line) != 2:
            continue
        src, dst = line
        dst = dst.split(' ', 1)[0]
        mapping[src] = dst
    ldd.wait()
    print(template_begin.strip())
    process(options.binary, mapping)
    for elf in mapping.values():
        if elf:
            process(elf, mapping)
    print(template_end.strip())

if __name__ == '__main__':
    main()

# vim:ts=4 sts=4 sw=4 et
