#!/usr/bin/env python3
import argparse
import glob
from nullroute.core import Core
import os
from pprint import pprint
import re
import subprocess
import sys

needed_re = re.compile(r'Shared library: \[(.+)\]$')
rpath_re = re.compile(r'Library r(?:un)?path: \[(.+)\]$')

def parse_ldconf(conf_path):
    if not os.path.exists(conf_path):
        print("warning: config file %s not found" % conf_path, file=sys.stderr)
        return
    for line in open(conf_path):
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        line = line.split()
        if line[0] == "include" and line[1]:
            for incl_path in glob.glob(line[1]):
                yield from parse_ldconf(incl_path)
        else:
            yield line[0]

def get_lib_paths():
    """
    Return an array of paths where ld.so searches for libraries.
    """
    paths = []
    if "LD_LIBRARY_PATH" in os.environ:
        paths += os.environ["LD_LIBRARY_PATH"].split(":")
    paths += parse_ldconf("/etc/ld.so.conf")
    paths += ["/lib", "/usr/lib"]
    paths.remove("/usr/lib32")
    return paths

def find_in_path(paths, basename):
    """
    Return the first existing path of a file `basename` in given paths.
    Return None if not found.

    >>> find_in_path(os.environ["PATH"], "sh")
    "/bin/sh"
    """
    if "/" in basename:
        return os.path.abspath(basename)
    for _dir in paths:
        path = os.path.join(_dir, basename)
        if os.path.exists(path):
            return path

def is_elf(path):
    """
    Check whether given path is an ELF file.
    """

def read_needed(path):
    """
    Read direct dependencies of an ELF file.
    """
    deps = set()
    rpaths = []
    proc = subprocess.Popen(["readelf", "-d", path],
                            stdout=subprocess.PIPE,
                            stderr=subprocess.DEVNULL)
    for line in proc.stdout:
        line = line.decode("utf-8")
        m = needed_re.search(line)
        if m:
            deps.add(m.group(1))
            continue
        m = rpath_re.search(line)
        if m:
            rpaths += m.group(1).split(":")
            # TODO: parse $ORIGIN, $LIB, relative rpath
            continue
    return deps, rpaths

# Functions for displaying dependency trees
# (will not work for "normal" trees or graphs)

def is_terminal():
    return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()

def show_tree(root, tree, indent=0, highlight=None, ctx=None):
    """
    Print dict `tree` {item: [children...]} starting at a given `root`
    as a textual tree. Recurse for each item in `tree[root]` as new root.

    >>> tree = {"a": {"b", "c"}, "b": {"d"}, "c": {"b", "d"}}
    >>> show_tree("a", tree)
    a
    ├─b
    │ └─d
    └─c
      ├─b
      │ └─d
      └─d
    """
    depth, branches, seen = ctx or (0, [], set())
    if depth == 0:
        print(" "*indent + root)
    if root not in tree:
        return
    branches += [None]
    seen |= {root}
    if not highlight:
        highlight = dict()
    children = set(tree[root]) - seen
    more = len(children)
    for child in sorted(children):
        more -= 1
        branches[depth] = ("├" if more else "└") + "─"
        prefix = suffix = ""
        if is_terminal() and child in highlight:
            color = highlight[child]
            prefix = "\033[%sm" % color
            suffix = "\033[m"
        if child in seen:
            continue
        print(" "*indent + "".join(branches) + prefix + child + suffix)
        if child in tree:
            branches[depth] = ("│" if more else " ") + " "
            ctx = depth + 1, branches.copy(), seen.copy()
            show_tree(child, tree, indent, highlight, ctx)

def walk_tree(prefix, tree):
    """
    Return a list containing all possible paths that start with `prefix`
    and exist in the `tree` dict. (`prefix` must be a list, even if it
    consists of a single item only.)

    >>> tree = {"a": {"b", "c"}, "b": {"c"}, "c": {"b", "d"}}
    >>> tree
    {'a': {'b', 'c'},
     'b': {'c'},
     'c': {'b', 'd'}}
    >>> list(walk_tree(["a"], tree))
    [['a', 'c', 'd'],
     ['a', 'c', 'b', 'd'],
     ['a', 'b', 'd']]
    """
    children = tree[prefix[-1]]
    for child in children:
        if child in prefix:
            raise ValueError("dependency loop detected at %r + %r" % (prefix, child))
        chain = prefix + [child]
        if child in tree and tree[child]:
            yield from walk_tree(chain, tree)
        else:
            yield chain

def flip_tree(roots, tree):
    """
    Return a dict of reverse dependencies where each root in `roots`
    becomes a leaf if formatted as a tree. (The output is usually
    formatted as *multiple* trees, one for each key in resulting tree.)

    >>> tree
    {'a': {'b', 'c'},
     'b': {'c'},
     'c': {'b', 'd'}}
    >>> flip_tree(["a"], tree)
    {'d': {'c', 'b'},
     'c': {'a'},
     'b': {'a', 'c'}}
    >>> show_tree("d", flip_tree(["a"], tree))
    d
    ├─b
    │ ├─a
    │ └─c
    │   └─a
    └─c
      └─a
    """
    flipped_tree = dict()
    for root in roots:
        for chain in walk_tree([root], tree):
            user = chain.pop()
            while chain:
                dep = chain.pop()
                if user not in flipped_tree:
                    flipped_tree[user] = set()
                flipped_tree[user].add(dep)
                user = dep
    return flipped_tree

class MissingLibScanner(object):
    def __init__(self, exe_path=None):
        self.exe_path = exe_path
        self.lib_paths = get_lib_paths()
        self.lib_rpaths = dict()
        self.resolved_paths = dict()
        self.forward_deps = dict()
        self.reverse_deps = dict()
        self.missing_libs = set()

    def find_missing(self):
        if not self.exe_path.startswith("/"):
            raise ValueError("exe_path must be an absolute path")

        self.missing_libs = set()

        todo = {(None, self.exe_path)}

        while todo:
            parent, elf_name = todo.pop()

            if elf_name in self.forward_deps:
                continue

            if elf_name in self.resolved_paths:
                elf_path = self.resolved_paths[elf_name]
            else:
                paths = self.lib_rpaths.get(parent, []) + self.lib_paths
                elf_path = find_in_path(paths, elf_name)
                self.resolved_paths[elf_name] = elf_path

            if elf_path is None:
                self.missing_libs.add(elf_name)
                continue

            deps, rpaths = read_needed(elf_path)
            self.forward_deps[elf_name] = deps
            self.lib_rpaths[elf_path] = rpaths

            for dep in deps:
                if dep not in self.reverse_deps:
                    self.reverse_deps[dep] = set()
                self.reverse_deps[dep].add(elf_name)

            todo |= {(elf_path, dep) for dep in deps}

        return self.missing_libs

    def show_reverse_tree(self):
        print("Reverse dependencies:")
        for lib in self.missing_libs:
            show_tree(lib, self.reverse_deps, indent=2)

    def show_forward_tree(self):
        flipped_deps = flip_tree(self.missing_libs, self.reverse_deps)
        print("Forward dependencies:")
        show_tree(self.exe_path, flipped_deps, indent=2,
              highlight={lib: "38;5;11" for lib in self.missing_libs})

parser = argparse.ArgumentParser()
parser.add_argument("file", nargs="+")
args = parser.parse_args()

os_paths = os.environ.get("PATH", "").split(":")
ldm = MissingLibScanner()

for exe_path in args.file:
    ldm.exe_path = find_in_path(os_paths, exe_path)
    #print("Scanning %s" % ldm.exe_path, file=sys.stderr)
    missing = ldm.find_missing()
    m = os.environ.get("PRETEND_MISSING", "")
    if m:
        missing |= set(m.split(","))
    if missing:
        ldm.show_forward_tree()
