#!/usr/bin/python
# coding: utf-8

# Copyright © Sébastien Luttringer
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

'''Find dependencies of a package or directory'''

from argparse import ArgumentParser
from elftools.elf.dynamic import DynamicSection, DynamicSegment
from elftools.elf.elffile import ELFFile
from os import walk, environ, getcwd, chdir, access, R_OK
from os.path import join, exists, isdir, isfile, normpath, realpath, splitext
from pprint import pformat
from pycman import config
from pygments import highlight
from pygments.formatters import TerminalFormatter, NullFormatter
from pygments.lexers import PythonLexer
from shlex import split
from sys import stdout, stderr
from tarfile import open as tar
from tempfile import TemporaryDirectory, NamedTemporaryFile

PACKAGES = None

def msg(message):
    '''display arch linux message'''
    if stdout.isatty():
        stdout.write("\033[1;32m==> \033[1;37m%s\033[m\n" % message)
    else:
        stdout.write("==> %s\n" % message)

def err(message):
    '''display colored error'''
    if stderr.isatty():
        stderr.write("\033[1;31mERROR: \033[1;37m%s\033[m\n" % message)
    else:
        stderr.write("ERROR: %s\n" % message)

def missing_pkginfo(path, deps):
    '''show deps against path'''
    fpath = join(path, ".PKGINFO")
    if not access(fpath, R_OK):
        return
    pkginfo = parse_pkginfo(fpath)
    if "depend" not in pkginfo:
        return
    return(set(deps.keys() - set(pkginfo["depend"])))

def parse_pkginfo(path):
    '''parse a .PKGINFO file'''
    pkginfo = dict()
    with open(path) as fd:
        for line in fd.readlines():
            # skip empty and comment
            if len(line) == 0 or line[0] == "#":
                continue
            lhs, rhs = line.split("=", 1)
            pkginfo.setdefault(lhs.strip(), []).append(rhs.strip())
    return(pkginfo)

def find_sharedlibs(fd):
    ef = ELFFile(fd)
    for section in ef.iter_sections():
        if isinstance(section, DynamicSection):
            for tag in section.iter_tags():
                if tag.entry.d_tag == 'DT_NEEDED':
                    yield(tag.needed)

def find_pkg(path):
    '''find the package owning path'''
    if not exists(path):
        return None
    path = normpath(realpath(path)).lstrip('/')
    global PACKAGES
    if PACKAGES is None:
        PACKAGES = config.init_with_config(environ.get("PACMAN_CONF",
            "/etc/pacman.conf")).get_localdb().pkgcache
    for pkg in PACKAGES:
        if path in [ x[0] for x in pkg.files]:
            return(pkg.name)
    return None

def find_deps(path):
    '''find deps packages'''
    deps = {}
    cwd = getcwd()
    chdir(path)
    for (dirpath, dirnames, filenames) in walk("."):
        for filename in filenames:
            fpath = join(dirpath, filename)
            with open(fpath, "rb") as fd:
                magic = fd.read(4)
                fd.seek(0)
                # ELF
                if magic == b"\x7fELF":
                    for libname in find_sharedlibs(fd):
                        pkgname = find_pkg("/usr/lib/%s" % libname)
                        deps.setdefault(pkgname, dict())
                        deps[pkgname].setdefault(libname, set())
                        deps[pkgname][libname] |= {fpath}
                elif magic[:2] == b"#!":
                    exec_path = fd.readline()[2:].split()[0].decode()
                    pkgname = find_pkg(exec_path)
                    deps.setdefault(pkgname, dict())
                    deps[pkgname].setdefault(exec_path, set())
                    deps[pkgname][exec_path] |= {fpath}
    chdir(cwd)
    return(deps)

def extract_zst(path):
    '''Extract zstandard file'''
    import zstandard as zstd
    wh = NamedTemporaryFile()
    with open(path, "rb") as rh:
        dctx = zstd.ZstdDecompressor()
        reader = dctx.stream_reader(rh)
        while True:
            chunk = reader.read(65536)
            if not chunk:
                break
            wh.write(chunk)
    wh.seek(0, 0)
    return wh

def parse_argv():
    '''Parse command line arguments'''
    parser = ArgumentParser()
    parser.add_argument("path", metavar="package|directory")
    return parser.parse_args()

def main():
    '''Program entry point'''
    args = parse_argv()
    if isfile(args.path):
        # tarball module doesn't yet support zst file, temporary workaround
        if splitext(args.path)[1] == ".zst":
            fo = extract_zst(args.path)
        else:
            fo = open(args.path, "rb")
        tarball = tar(fileobj=fo)
        pkgdir = TemporaryDirectory()
        tarball.extractall(pkgdir.name)
        args.path = pkgdir.name
    if isdir(args.path):
        deps = find_deps(args.path)
        mdeps = missing_pkginfo(args.path, deps)
        # display deps
        msg("Found deps")
        formater = TerminalFormatter() if stdout.isatty() else NullFormatter()
        stdout.write(highlight(pformat(deps), PythonLexer(), formater))
        # show missing deps in .PKGINFO
        if mdeps is None:
            msg(".PKGINFO not found")
        elif len(mdeps) > 0:
            msg("Missing in .PKGINFO")
            stdout.write(highlight(pformat(mdeps), PythonLexer(), formater))
        else:
            msg("Nothing is missing in .PKGINFO")
    else:
        err("Unsupported file type")
        return 2
    return 0

if __name__ == '__main__':
    main()

# vim:set ts=4 sw=4 et ai: