LLDB screengraph

lldb_screengraph.py

#!/usr/bin/python

# ---------------------------------------------------------------------
# Be sure to add the python path that points to the LLDB shared library.
#
# # To use this in the embedded python interpreter using "lldb" just
# import it with the full path using the "command script import"
# command
#   (lldb) command script import /path/to/lldb_screengraph.py
# ---------------------------------------------------------------------

from __future__ import print_function

import inspect
import lldb
import optparse
import os
import shlex
import sys
import textwrap

from sets import Set


def debug(s):
    print(s)

def make_directory_if_not_exist(directory):
    try:
        os.makedirs(directory)
    except OSError, e:
        if e.errno != os.errno.EEXIST:
            raise

class Output:
    
    def append(self, state):
        raise Exception('not implemented')

class TextOutput(Output):
    
    def __init__(self, directory):
        self.filename = os.path.join(directory, 'trace.txt')
        self.states = []
        
    def append(self, state):
        self.states.append(state)
        self.write()

    def write(self):
        output = ''
        for state in self.states:
            output = output + str(state)
        debug(output)
        with open(self.filename, 'w+') as f:
            f.write(output)
            
class ScreenshotOutput(Output):
    
    def __init__(self, directory):
        self.directory = directory
        self.count = 0
    
    def append(self, state):
        self.screenshot(state)
    
    @property
    def filename(self):
        return os.path.join(self.directory, 'screenshot%i.png' % self.count)
        
    def screenshot(self, state):
        options = lldb.SBExpressionOptions()
        options.SetLanguage(lldb.eLanguageTypeSwift)
        s = """
        UIGraphicsBeginImageContext(view.bounds.size)
        if let view = UIApplication.shared.keyWindow,
            let ctx = UIGraphicsGetCurrentContext(),
            let fileURL = URL(string: "file://%s") {
            view.layer.render(in: ctx);
            if let image = UIGraphicsGetImageFromCurrentImageContext(),
                let data = image.pngData() {
                do {
                    try data.write(to: fileURL)
                    print("Screenshot saved")
                } catch {
                    print("Error saving screenshot:", error)
                }
            }
        }
        UIGraphicsEndImageContext()
        """ % (self.filename)
    
        value = state.frame.EvaluateExpression(s, options)
        error = value.GetError()
    
        if error.description:
            print('Error saving screenshot: ' + error.description)
        else:
            self.count = self.count + 1


class GraphvizOutput(Output):
    
    class Edge:
        def __init__(self, src, dst):
            self.src = src
            self.dst = dst
        
        def __str__(self):
            # return 'N%i -> N%i;' % (self.src.index, self.dst.index)
            return 'N%i -> N%i [ label = "%s" ];' % (
                self.src.index, 
                self.dst.index,
                str(self.src.state),
            )
            

    class Node:
        def __init__(self, state, index=0):
            self.state = state
            self.edges = []
            self.index = index

        def add_edge(self, node):
            self.edges.append(GraphvizOutput.Edge(self, node))
        
        def __str__(self):
            return 'N%i [shape=rect,image="screenshot%i.png", label="%s", labelloc=b];' % (
                self.index,
                self.index,
                #str(self.state),
                '',
            )

    def __init__(self, directory):
        self.filename = os.path.join(directory, 'graph.dot')
        self.root = None
        self.last = None
        self.nodes = dict()
        self.count = 0
        
    def append(self, state):
        key = state.breakpoint_id
    
        if key in self.nodes.keys():
            node = self.nodes[key]
        else:
            node = GraphvizOutput.Node(state, self.count)
            self.nodes[key] = node
            self.count = self.count + 1
        
        if not self.root:
            self.root = node
        
        if self.last:
            self.last.add_edge(node)
        
        self.last = node
        
        with open(self.filename, 'w+') as f:
            f.write(self.output)

    @property
    def output(self): #TODO rewrite this to reduce complexity (like https://stackoverflow.com/a/10289740)
        queue = [self.root]
        visited = Set()
        
        nodes = ''
        edges = ''
        
        while queue:
            node = queue.pop(0)
            nodes = nodes + str(node)
            visited.add(node)
            for edge in node.edges:
                edges = edges + str(edge)
                if not edge.dst in visited:
                    queue.append(edge.dst)
            
        output = textwrap.dedent('''
            digraph G {
                rankdir = LR;
                %s
                %s
            }
        ''') % (nodes, edges)
        return output


class State:
    
    class Variable:
        
        def __init__(self, variable):
            self.name = variable.name
            self.type = variable.type.name
            self.info = variable.description
        
        def __repr__(self):
            return '<Variable %s: %s>' % (self.name, self.type)
            
        def __str__(self):
            return '%s: %s (%s)' % (self.name, self.type, self.info)
    
    def __init__(self, frame, location):
        
        self.frame = frame
        self.location = location

        breakpoint = location.GetBreakpoint()
        self.breakpoint_id = str(breakpoint.id)
        self.function = location.GetAddress().function.name
        self.line_entry = '%s:%i' % (
            str(location.GetAddress().line_entry.GetFileSpec()),
            location.GetAddress().line_entry.GetLine(),
        )

        self.variables = []
        variables_list = frame.GetVariables(True, False, False, False)
        variables_count = variables_list.GetSize()
        for i in range(0, variables_count):
            variable = variables_list.GetValueAtIndex(i)
            self.variables.append(State.Variable(variable))
            
    def __repr__(self):
        return '<State: breakpoint %s (%s), %i variables>' % (
            self.breakpoint_id,
            self.line_entry,
            len(self.variables),
        )
    
    def __str__(self):
        description = textwrap.dedent('''
            breakpoint = %s
                function = %s
                line_entry = %s
        ''') % (
            self.breakpoint_id,
            self.function,
            self.line_entry,
        )
        for variable in self.variables:
            description = description + '\tvariable = %s' % str(variable)
        return description        


class TracingSession:

    def __init__(self, debugger, outputs):
        TracingSession.current = self
        self.debugger = debugger
        self.outputs = outputs
            
    def start(self):
        print('starting screengraph')
        self.breakpoints = []
        target = self.debugger.GetSelectedTarget()
        for breakpoint in target.breakpoint_iter():
            if breakpoint.IsValid() and breakpoint.IsEnabled():
                debug(breakpoint)
                breakpoint.SetScriptCallbackFunction('screengraph.TracingSession.on_breakpoint_hit')
                self.breakpoints.append(breakpoint)
    
    def stop(self):
        print('stopping screengraph')
        for breakpoint in self.breakpoints:
            if breakpoint.IsValid():
                dir(breakpoint)
                breakpoint.SetScriptCallbackFunction("")
        self.breakpoints = []
        
    @staticmethod
    def on_breakpoint_hit(frame, location, internal_dict):
        if frame.IsValid():
            debug('Hit frame: ' + str(frame))
            TracingSession.current.process(frame, location, internal_dict)
    
    def process(self, frame, location, internal_dict):
        state = State(frame, location)
        debug(repr(state))
        for output in self.outputs:
            output.append(state)
        frame.GetThread().GetProcess().Continue()
    

class ScreenGraphCommand:
    program = 'screengraph'

    @classmethod
    def register_lldb_command(cls, debugger, module_name):
        parser = cls.create_options()
        cls.__doc__ = parser.format_help()
        command = 'command script add -c %s.%s %s' % (module_name,
                                                      cls.__name__,
                                                      cls.program)
        debugger.HandleCommand(command)
        print('The "{0}" command has been installed, type "help {0}" or "{0} '
              '--help" for detailed help.'.format(cls.program))

    @classmethod
    def create_options(cls):
        usage = "usage: %prog start|stop"
        description = ('Creates a graph of screens.')
        
        parser = optparse.OptionParser(
            description=description,
            prog=cls.program,
            usage=usage,
            add_help_option=False)
        
        return parser
        
    def get_short_help(self):
        return 'Creates a graph of screens.'

    def get_long_help(self):
        return self.self.parser.format_help()

    def __init__(self, debugger, unused):
        self.parser = self.create_options()
        self.tracing = None
        
    def __call__(self, debugger, command, exe_ctx, result):
        command_args = shlex.split(command)
        try:
            (options, args) = self.parser.parse_args(command_args)
        except:
            result.SetError("Option parsing failed")
            return

        if not args:
            self.parser.print_help()
            return
        subcommand = args[0]
        
        if subcommand == 'start':
            directory = os.path.join(os.path.expanduser("~"), 'screengraph')
            
            outputs = self.outputs(directory)
            
            tracing = TracingSession(debugger, outputs)
            tracing.start()
            self.tracing = tracing
            
        elif subcommand == 'stop':
            if self.tracing:
                self.tracing.stop()
                
    def outputs(self, directory, text=True, screenshot=True, graphviz=True):
        make_directory_if_not_exist(directory)
            
        outputs = []
        if text:
            outputs.append(TextOutput(directory))
        if screenshot:
            outputs.append(ScreenshotOutput(directory))
        if graphviz:
            outputs.append(GraphvizOutput(directory))
        return outputs

def __lldb_init_module(debugger, dict):   
    for _name, cls in inspect.getmembers(sys.modules[__name__]):
        if inspect.isclass(cls) and callable(getattr(cls,
                                                     "register_lldb_command",
                                                     None)):
            cls.register_lldb_command(debugger, __name__)