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

# Copyright 2020 xaizek <xaizek@posteo.net>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os

EXPORT_EDGE = "edge [style=filled,weight=1];"
IMPORT_EDGE = "edge [style=dashed];"

INVALID = -1
EXPORT = 1
IMPORT = 2
OBJECT_EXPORT = 3
OBJECT_IMPORT = 4

symbols = {}
objects = []

class objModule:
    def __init__(self, fileName, name = None):
        self._fileName = fileName
        self._name = os.path.basename(fileName) if name is None else name
        self._importTable = []
        self._exportTable = []
    def addImport(self, entry):
        self._importTable += [entry]
        entry.addImport(self)
    def remImport(self, entry):
        self._importTable.remove(entry)
    def addExport(self, entry):
        self._exportTable += [entry]
        entry.addExport(self)
    def remExport(self, entry):
        self._exportTable.remove(entry)
    def interactsWith(self, mod):
        for entry in self._importTable:
            if entry.belongsTo(mod):
                return True
        return False
    def bindsCount(self, mod):
        count = 0
        for entry in self._importTable:
            if entry.belongsTo(mod):
                count += 1
        return count
    def getImports(self):
        return [x for x in self._importTable if x.isUsers()]
    def getExports(self):
        return [x for x in self._exportTable if x.isUsers()]
    def getAllExports(self):
        return self._exportTable
    def printMe(self):
        if len(self.getImports()) == 0 and len(self.getExports()) == 0:
            return

        if len(self.getImports()) == 0:
            color = 'green'
        elif len(self.getExports()) == 0:
            color = 'burlywood'
        else:
            color = 'yellow'
        print '"%s" [shape=rectangle,style=filled,fillcolor=%s];' \
                % (self._name, color)
    def getName(self):
        return self._name
    def __str__(self):
        return self._name
    def getFileName(self):
        return self._fileName

class function:
    def __init__(self, name):
        self.name = name
        self.importedBy = []
        self._exportCounter = 0
        self._importCounter = 0
        self._selfImportCounter = 0
        self._mod = None
    def remove(self):
        for mod in self.importedBy:
            mod.remImport(self)
        self._mod.remExport(self)
    def addImport(self, mod):
        if mod == self:
            self._selfImportCounter += 1
        else:
            self._importCounter += 1
        self.importedBy += [mod]
    def addExport(self, mod):
        assert self._mod is None, ("%s: %s & %s" %
               (self.name, self._mod.getName(), mod.getName()))
        self._exportCounter = 1
        self._mod = mod;
    def isUsers(self):
        return self._exportCounter > 0 and self._importCounter > 0
    def isSystem(self):
        return self._exportCounter == 0 and self._importCounter > 0
    def isUnused(self):
        return self._exportCounter > 0 \
           and self._importCounter - self._selfImportCounter == 0
    def belongsTo(self, mod):
        return self._mod == mod
    def getMod(self):
        return self._mod
    def __str__(self):
        return self.name

def getFuncObjdump(line):
    parts = line.split()
    if len(parts) == 6 and parts[1] == 'g' and parts[2] == 'F':
        return (EXPORT, parts[5])
    elif len(parts) == 6 and parts[1] == 'g' and parts[2] == 'O':
        return (OBJECT_EXPORT, parts[5])
    elif len(parts) == 5 and parts[1] == 'O' and parts[2] == '*COM*':
        return (OBJECT_IMPORT, parts[4])
    elif len(parts) == 4 and parts[1] == '*UND*':
        return (IMPORT, parts[3])
    else:
        return (INVALID, None)

def getTables(objFile):
    global objects
    with os.popen('objdump -t ' + objFile.getFileName()) as f:
        lines = f.readlines()
        for line in lines:
            (lineType, name) = getFuncObjdump(line)
            if lineType == INVALID:
                continue
            if lineType == OBJECT_EXPORT or lineType == OBJECT_IMPORT:
                objects += [name]

            if symbols.has_key(name):
                func = symbols[name]
            else:
                func = function(name)
                symbols[name] = func

            if lineType == IMPORT or lineType == OBJECT_IMPORT:
                objFile.addImport(func)
            elif lineType == EXPORT or lineType == OBJECT_EXPORT:
                objFile.addExport(func)

def printFunctions(mod, funcs):
    for entry in funcs:
        print '"%s"->"%s";' % (mod, entry)

def printSymGraph(objModules):
    for mod in objModules:
        mod.printMe()

        # exports
        print EXPORT_EDGE
        printFunctions(mod, mod.getExports())

        # imports
        print IMPORT_EDGE
        printFunctions(mod, mod.getImports())

def printOneModFuncGraph(mod):
    mod.printMe()

    modImports = mod.getImports()

    # imports
    print IMPORT_EDGE
    for entry in modImports:
        print '"%s"->"%s";' % (mod, entry)

    # exports
    print EXPORT_EDGE
    for entry in modImports:
        entry.getMod().printMe()
        print '"%s"->"%s";' % (entry.getMod(), entry)

def printOneInvModFuncGraph(mod):
    mod.printMe()

    modExports = mod.getExports()

    # exports
    print EXPORT_EDGE
    for entry in modExports:
        print '"%s"->"%s";' % (mod, entry)

    # imports
    print IMPORT_EDGE
    for entry in modExports:
        for mod in entry.importedBy:
            mod.printMe()
            print '"%s"->"%s";' % (mod, entry)

def analyzeMods(x, y):
    if not x.interactsWith(y):
        return

    x.printMe()
    y.printMe()
    print 'edge [label="%d"];' % x.bindsCount(y)
    print '"%s"->"%s";' % (x, y)

def printModGraph(objModules):
    [analyzeMods(x, y) for x in objModules for y in objModules]

def printSingleModGraph(mod, modules):
    for x in modules:
        if x != mod:
            analyzeMods(x, mod)

def printOneInvModGraph(mod, modules):
    for x in modules:
        if x != mod:
            analyzeMods(mod, x)

def printSymbolGraph(func):
    func.getMod().printMe()

    # export
    print EXPORT_EDGE
    print '"%s"->"%s";' % (func.getMod(), func)

    # imports
    print IMPORT_EDGE
    for mod in func.importedBy:
        mod.printMe()
        print '"%s"->"%s";' % (mod, func)

def makeModule(objFileName):
    objFile = objModule(objFileName)
    getTables(objFile)
    return objFile

def printLegend():
    print '"export only" [shape=rectangle,style=filled,fillcolor=green];'
    print '"import only" [shape=rectangle,style=filled,fillcolor=burlywood];'
    print '"export and import" [shape=rectangle,style=filled,fillcolor=yellow];'

def listUnusedFuncs(mods):
    for mod in mods:
        funcs = [x for x in mod.getAllExports() if x.isUnused()]
        if len(funcs) == 0:
            continue

        print '--- unused from ' + str(mod)
        for func in funcs:
            print func

def listSystemFuncs():
    funcNames = [x.name for x in symbols.values() if x.isSystem()]
    funcNames.sort()
    for funcName in funcNames:
        print funcName

def listUserFuncs():
    userFuncs = [x for x in symbols.values() if x.isUsers() or x.isUnused()]
    userFuncs = sorted(userFuncs, None, lambda x: x.name)
    print 'Name ExportedFrom ImportedBy'
    for func in userFuncs:
        lst = ', '.join(sorted(map(lambda x: str(x), func.importedBy)))
        print '%s %s %s (%s)' % (func.name, func.getMod(), len(func.importedBy),
                lst)

def dirGroupName(path, root):
    if not root:
        return os.path.dirname(path)
    return os.path.relpath(os.path.dirname(path), root)

def main():
    parser = argparse.ArgumentParser(description='Object file analyzer')
    parser.add_argument('--moddeps', dest='moddeps', action='store_true',
                        help='Generate module dependency graph')
    parser.add_argument('--symdeps', dest='symdeps', action='store_true',
                        help='Generate symbol dependency graph')
    parser.add_argument('--module', dest='module',
                        help='Generate graph for this module (its users)')
    parser.add_argument('--invert', dest='invert', action='store_true',
                        help='Invert dependencies for --module')
    parser.add_argument('--symbols', dest='symbols', action='store_true',
                        help='Detail symbols used by --module')
    parser.add_argument('--objects', dest='objects', action='store_true',
                        help='Display information on objects')
    parser.add_argument('--symbol', dest='symbol',
                        help='Generate graph for this symbol')
    parser.add_argument('--system', dest='system', action='store_true',
                        help='List references to system symbols')
    parser.add_argument('--user', dest='user', action='store_true',
                        help='List table of exported user symbols')
    parser.add_argument('--unused', dest='unused', action='store_true',
                        help='List exported, but unused symbols')
    parser.add_argument('--group-dirs', dest='group_dirs', action='store_true',
                        help='Group files inside directories into units')
    parser.add_argument('--root', dest='root',
                        help='Group files inside directories into units')
    parser.add_argument('obj', nargs='+',
                        help='Object file to get information from')
    args = parser.parse_args()

    if not args.unused and not args.system and not args.user and \
       not args.symbol and not args.moddeps and not args.symdeps and \
       not args.module:
        print 'No action selected'
        exit(1)

    # parse object files
    modulesList = [makeModule(x) for x in args.obj]
    objModules = {}
    for mod in modulesList:
        objModules[mod.getName()] = mod

    if args.group_dirs:
        newObjModules = {}
        newSymbols = {}
        for mod in modulesList:
            dirName = dirGroupName(mod.getFileName(), args.root)
            newObjModules[dirName] = objModule(dirName, dirName + "/*")
        for modA in modulesList:
            modAdirName = dirGroupName(modA.getFileName(), args.root)
            for export in modA.getAllExports():
                name = export.name
                if newSymbols.has_key(name):
                    func = newSymbols[name]
                else:
                    func = function(name)
                    newSymbols[name] = func
                    newObjModules[modAdirName].addExport(func)

                for modB in export.importedBy:
                    modBdirName = dirGroupName(modB.getFileName(), args.root)
                    newObjModules[modBdirName].addImport(func)

        global symbols
        global objects
        symbols = newSymbols
        objects = []

        objModules = newObjModules

    # add special system module
    sysMod = objModule('system')
    objModules['system'] = sysMod
    for sym in symbols.values():
        if sym.isSystem():
            sysMod.addExport(sym)

    if 'main' in symbols:
        sysMod.addImport(symbols['main'])

    if args.objects:
        # remove non-objects from list of symbols
        for symname in symbols.keys():
            if symname not in objects:
                symbols[symname].remove()
                del symbols[symname]
    else:
        # remove objects from list of symbols
        for obj in objects:
            if obj not in symbols:
                symbols[obj].remove()
                del symbols[obj]

    if args.symdeps:
        print 'digraph "symdeps" {'
        printSymGraph(objModules.values())
        print '}'

    if args.moddeps:
        print 'digraph "moddeps" {'
        printModGraph(objModules.values())
        print '}'

    if args.module:
        if args.symbols:
            if args.invert:
                print 'digraph "singlemodfuncdeps" {'
                printOneInvModFuncGraph(objModules[args.module])
                print '}'
            else:
                print 'digraph "singlemodfuncusers" {'
                printOneModFuncGraph(objModules[args.module])
                print '}'
        else:
            if args.invert:
                print 'digraph "singlemoddeps" {'
                printOneInvModGraph(objModules[args.module],
                                    objModules.values())
                print '}'
            else:
                print 'digraph "singlemodusers" {'
                printSingleModGraph(objModules[args.module],
                                    objModules.values())
                print '}'

    if args.symbol:
        print 'digraph "onesymdeps" {'
        printSymbolGraph(symbols[args.symbol])
        print '}'

    if args.unused:
        listUnusedFuncs(objModules.values())

    if args.system:
        listSystemFuncs()

    if args.user:
        listUserFuncs()

main()