Updating memory drawer to allow both files and pd dataframes

debug.py

@@ -1,11 +1,18 @@
+"""
+Sample file for debugging purposes and examples
+"""
+
+import pandas as pd
+
 #from herrewebpy.bioinformatics import sequence_alignment
 #sequence_alignment.SequenceAlignment(['aa', 'bb', 'cc'],['bb','aa','cc'], ['1','2','3'], ['1','2','3'])
 
 #from herrewebpy.firmware_forensics import function_extractor
 #function_extractor.FunctionExtractor('', 'ARM_AARCH64')
 
-#from herrewebpy.christianity import readplan_generator
-#readplan_generator.generate_readplan()
+# from herrewebpy.christianity import readplan_generator
+# readplan_generator.generate_readplan()
 
-from herrewebpy.firmware_forensics import memory_drawer
-memory_drawer.MemoryDrawer('sample_data/csv/stack_and_functions.csv')
+from herrewebpy.firmware_forensics.memory_drawer import MemoryDrawer
+df = pd.read_csv('sample_data/csv/stack_and_functions.csv')
+MemoryDrawer(df)

herrewebpy/christianity/readplan_generator.py

@@ -30,7 +30,12 @@ def generate_readplan(start_date):
     total_chapters = sum([bible.get_number_of_chapters(reading_list[i]) for i in range(len(reading_list))])
     chapters_per_day = total_chapters // 365 + 1
 
-    df = pd.DataFrame(columns=['Book', 'Chapters'])
+    # Create a dataframe with each book, each chapter, and number of verses
+    df = pd.DataFrame(columns=['Book', 'Chapters', 'Verses'])
+    for book in reading_list:
+        df = pd.concat([df, pd.DataFrame({'Book': [book.title], 'Chapters': [bible.get_number_of_chapters(book)]})])
+
+    df = pd.DataFrame(columns=['Book', 'Chapters', 'Verses'])
     for book in reading_list:
         df = pd.concat([df, pd.DataFrame({'Book': [book.title], 'Chapters': [bible.get_number_of_chapters(book)]})])
 

herrewebpy/firmware_forensics/memory_drawer.py

@@ -1,7 +1,6 @@
 # Using plotly
 import plotly.graph_objects as go
-import random, argparse
-import numpy as np
+import random, argparse, os, datetime
 import pandas as pd
 
 """
@@ -14,261 +13,283 @@ This script reads a CSV file with the following columns: start,end,name,order,co
 Then it generates a memory map of the regions, and outputs an HTML file with the memory map.
 """
 
-def read_data(input_file):
-    data = pd.read_csv(input_file)
+class MemoryDrawer():
 
-    def convert_to_int(value):
-        try:
-            if isinstance(value, str) and value.startswith('0x'):
-                return int(value, 16)
+    def __init__(self, input):
+        """
+        If this file is run manually, will take an input .csv path and output a memory map in .html format.
+
+        Args:
+            (Required) input (str): Path to the input .csv file
+            (Optional) output (str): Path to the output .html file
+        """
+        if isinstance(input, str):
+            if os.path.isfile(input):
+                output = f'{os.path.splitext(os.path.basename(input))[0]}_memory_drawer'
+                data = MemoryDrawer.read_data(pd.read_csv(input))
             else:
-                return int(value)
-        except ValueError:
-            return value 
-
-    data['start'] = data['start'].apply(convert_to_int)
-    data['end'] = data['end'].apply(convert_to_int)
-    data['size'] = data['end'] - data['start']
-
-    #data.sort_values(by=['size'], inplace=True, ascending=False)
-    data.sort_values(by=['start', 'size'], inplace=True, ascending=True)
-
-    # Inverse the order of the data
-    data.reset_index(drop=True, inplace=True)
-
-    data['overlap'] = False
-    data['index'] = data.index
-
-    for i, row in data.iterrows():
-        data.at[i, 'overlap'] = False
-        data.at[i, 'partial_overlap'] = False
-
-        # Annotate rows that fully overlap the current row
-        temp = data.loc[(data['start'] <= row['start']) & (data['end'] >= row['end'])]
-        if temp.shape[0] > 1:
-            data.at[i, 'overlap'] = True
-            data.at[i, 'overlapped_by'] = ','.join(temp['index'].astype(str).to_list())
-
-        # Annotate rows that partially overlap the current row (from start, but not to end)
-        temp = data.loc[(data['start'] <= row['start']) & (data['end'] < row['end']) & (data['end'] >= row['start'])]
-        if temp.shape[0] > 1:
-            data.at[i, 'partial_overlap'] = "Bottom"
-            data.at[i, 'partial_overlapped_by'] = ','.join(temp['index'].astype(str).to_list())
-
-        # Annotate rows that partially overlap the current row (from end, but not to start)
-        temp = data.loc[(data['start'] > row['start']) & (data['end'] >= row['end']) & (data['start'] <= row['end'])]
-        if temp.shape[0] > 1:
-            data.at[i, 'partial_overlap'] = "Top"
-            data.at[i, 'partial_overlapped_by'] = ','.join(temp['index'].astype(str).to_list())
-
-        # Also annotate which regions this row is overlapping
-        temp = data.loc[(data['start'] >= row['start']) & (data['end'] <= row['end'])]
-        if temp.shape[0] > 1:
-            data.at[i, 'overlap'] = True
-            data.at[i, 'overlapping'] = ','.join(temp['index'].astype(str).to_list())
-
-    # Send warnings if sizes are negative
-    if (data['size'] < 0).any():
-        print(f'Warning: Negative sizes detected at indices {data[data["size"] < 0].index}')
-    
-    return data
+                raise ValueError('Input string must be a path to a .csv file')
+        elif isinstance(input, pd.DataFrame):
+            now = datetime.datetime.now()
+            output = f'{now.strftime("%Y-%m-%d_%H-%M-%S")}_memory_drawer'
+            data = MemoryDrawer.read_data(input)
+        else:
+            raise ValueError('Input must be a path to a .csv file or a pandas DataFrame')
 
 
-def draw_diagram(data, vertical_gap_percentage=0.08, horizontal_gap=0.1):
-    tickpointers = []
-    labels = pd.DataFrame()
+        fig = MemoryDrawer.draw_diagram(data)
+        MemoryDrawer.write_output(fig, output)
 
-    def random_color():
-        return f'#{random.randint(0, 0xFFFFFF):06x}'
 
-    fig = go.Figure()
-    fig.update_layout(font=dict(family="Courier New, monospace"))
+    def read_data(data):
 
-    fig.update_layout(
-        plot_bgcolor='#FFFFFF',
-    )
+        def _convert_to_int(value):
+            try:
+                if isinstance(value, str) and value.startswith('0x'):
+                    return int(value, 16)
+                else:
+                    return int(value)
+            except ValueError:
+                return value 
 
-    for i, d in data.iterrows():
-        fillcolor = random_color()
-        data.at[i, 'fillcolor'] = fillcolor
+        data['start'] = data['start'].apply(_convert_to_int)
+        data['end'] = data['end'].apply(_convert_to_int)
+        data['size'] = data['end'] - data['start']
+
+        #data.sort_values(by=['size'], inplace=True, ascending=False)
+        data.sort_values(by=['start', 'size'], inplace=True, ascending=True)
+
+        # Inverse the order of the data
+        data.reset_index(drop=True, inplace=True)
+
+        data['overlap'] = False
+        data['index'] = data.index
+
+        for i, row in data.iterrows():
+            data.at[i, 'overlap'] = False
+            data.at[i, 'partial_overlap'] = False
+
+            # Annotate rows that fully overlap the current row
+            temp = data.loc[(data['start'] <= row['start']) & (data['end'] >= row['end'])]
+            if temp.shape[0] > 1:
+                data.at[i, 'overlap'] = True
+                data.at[i, 'overlapped_by'] = ','.join(temp['index'].astype(str).to_list())
+
+            # Annotate rows that partially overlap the current row (from start, but not to end)
+            temp = data.loc[(data['start'] <= row['start']) & (data['end'] < row['end']) & (data['end'] >= row['start'])]
+            if temp.shape[0] > 1:
+                data.at[i, 'partial_overlap'] = "Bottom"
+                data.at[i, 'partial_overlapped_by'] = ','.join(temp['index'].astype(str).to_list())
+
+            # Annotate rows that partially overlap the current row (from end, but not to start)
+            temp = data.loc[(data['start'] > row['start']) & (data['end'] >= row['end']) & (data['start'] <= row['end'])]
+            if temp.shape[0] > 1:
+                data.at[i, 'partial_overlap'] = "Top"
+                data.at[i, 'partial_overlapped_by'] = ','.join(temp['index'].astype(str).to_list())
+
+            # Also annotate which regions this row is overlapping
+            temp = data.loc[(data['start'] >= row['start']) & (data['end'] <= row['end'])]
+            if temp.shape[0] > 1:
+                data.at[i, 'overlap'] = True
+                data.at[i, 'overlapping'] = ','.join(temp['index'].astype(str).to_list())
+
+        # Send warnings if sizes are negative
+        if (data['size'] < 0).any():
+            print(f'Warning: Negative sizes detected at indices {data[data["size"] < 0].index}')
         
-        # Set base x values. Width of the rectangle.
-        x0 = 1
-        x1 = 6
+        return data
 
-        # Set base y values. Height of the rectangle.
-        y0 = d['index']
-        y1 = d['index']+1
 
-        if d['overlap'] == True:
-            # Row is overlapping the current row
-            if pd.notna(d['overlapping']):
-                y0 = sorted(map(int, d['overlapping'].split(',')))[0]
-                y1 = sorted(map(int, d['overlapping'].split(',')))[-1] + 1
+    def draw_diagram(data, vertical_gap_percentage=0.08, horizontal_gap=0.1):
+        tickpointers = []
+        labels = pd.DataFrame()
 
-            if pd.notna(d['overlapped_by']):
-                y0 = y0 + vertical_gap_percentage
-                y1 = y1 - vertical_gap_percentage
-                x0 = x0 + horizontal_gap
-                x1 = x1 - horizontal_gap
+        def random_color():
+            return f'#{random.randint(0, 0xFFFFFF):06x}'
 
-            if d['partial_overlap'] == "Bottom":
-                if pd.notna(d['partial_overlapped_by']):
-                    y0 = y0 + 0.25 + (0.6**len(d['partial_overlapped_by'].split(',')))
-                    #x0 = x0 + horizontal_gap
-                    #x1 = x1 - horizontal_gap
+        fig = go.Figure()
+        fig.update_layout(font=dict(family="Courier New, monospace"))
 
-            if d['partial_overlap'] == "Top":
-                if pd.notna(d['partial_overlapped_by']):
-                    y1 = y1 - (0.6**len(d['partial_overlapped_by'].split(',')))
-                    #x0 = x0 + horizontal_gap
-                    #x1 = x1 - horizontal_gap
-
-        fig.add_shape(
-            type="rect",
-            x0=x0,
-            x1=x1,
-            y0=y0+vertical_gap_percentage,
-            y1=y1-vertical_gap_percentage,
-            line=dict(width=1),
-            fillcolor=fillcolor,
-            opacity=0.4,
-            layer="below",
+        fig.update_layout(
+            plot_bgcolor='#FFFFFF',
         )
 
-        ### Add middle text
-        fig.add_trace(go.Scatter
-        (
-            x=[(x0+x1)/2],
-            y=[i+0.5],
-            text=d['name'],
-            mode="text",
-            textposition="middle center",
-            name=d['name'],
-            marker=dict(
-                color=fillcolor,
+        for i, d in data.iterrows():
+            fillcolor = random_color()
+            data.at[i, 'fillcolor'] = fillcolor
+            
+            # Set base x values. Width of the rectangle.
+            x0 = 1
+            x1 = 6
+
+            # Set base y values. Height of the rectangle.
+            y0 = d['index']
+            y1 = d['index']+1
+
+            if d['overlap'] == True:
+                # Row is overlapping the current row
+                if pd.notna(d['overlapping']):
+                    y0 = sorted(map(int, d['overlapping'].split(',')))[0]
+                    y1 = sorted(map(int, d['overlapping'].split(',')))[-1] + 1
+
+                if pd.notna(d['overlapped_by']):
+                    y0 = y0 + vertical_gap_percentage
+                    y1 = y1 - vertical_gap_percentage
+                    x0 = x0 + horizontal_gap
+                    x1 = x1 - horizontal_gap
+
+                if d['partial_overlap'] == "Bottom":
+                    if pd.notna(d['partial_overlapped_by']):
+                        y0 = y0 + 0.25 + (0.6**len(d['partial_overlapped_by'].split(',')))
+                        #x0 = x0 + horizontal_gap
+                        #x1 = x1 - horizontal_gap
+
+                if d['partial_overlap'] == "Top":
+                    if pd.notna(d['partial_overlapped_by']):
+                        y1 = y1 - (0.6**len(d['partial_overlapped_by'].split(',')))
+                        #x0 = x0 + horizontal_gap
+                        #x1 = x1 - horizontal_gap
+
+            fig.add_shape(
+                type="rect",
+                x0=x0,
+                x1=x1,
+                y0=y0+vertical_gap_percentage,
+                y1=y1-vertical_gap_percentage,
+                line=dict(width=1),
+                fillcolor=fillcolor,
+                opacity=0.4,
+                layer="below",
+            )
+
+            ### Add middle text
+            fig.add_trace(go.Scatter
+            (
+                x=[(x0+x1)/2],
+                y=[i+0.5],
+                text=d['name'],
+                mode="text",
+                textposition="middle center",
+                name=d['name'],
+                marker=dict(
+                    color=fillcolor,
+                ),
+            ))
+
+            ### Add top-left text with d['end']
+            # Overlapped to the right, to make it more readable
+            if pd.notna(d['overlapped_by']):
+                fig.add_trace(go.Scatter
+                (
+                    x=[(x1-0.24+horizontal_gap)],
+                    y=[y1-0.16],
+                    text=hex(d['end']),
+                    mode="text",
+                    textposition="middle center",
+                    marker=dict(
+                        color=fillcolor,
+                    ),
+                    showlegend=False,
+                ))
+
+                # Add bottom-left text with d['end']
+                fig.add_trace(go.Scatter
+                (
+                    x=[(x1-0.24+horizontal_gap)],
+                    y=[y0+0.14],
+                    text=hex(d['start']),
+                    mode="text",
+                    textposition="middle center",
+                    marker=dict(
+                        color=fillcolor,
+                    ),
+                    showlegend=False,
+                ))
+            else:
+                fig.add_trace(go.Scatter
+                (
+                    x=[(x0+0.14+horizontal_gap)],
+                    y=[y1-0.16],
+                    text=hex(d['end']),
+                    mode="text",
+                    textposition="middle center",
+                    marker=dict(
+                        color=fillcolor,
+                    ),
+                    showlegend=False,
+                ))
+
+                ### Add bottom-left text with d['end']
+                fig.add_trace(go.Scatter
+                (
+                    x=[(x0+0.14+horizontal_gap)],
+                    y=[y0+0.14],
+                    text=hex(d['start']),
+                    mode="text",
+                    textposition="middle center",
+                    marker=dict(
+                        color=fillcolor,
+                    ),
+                    showlegend=False,
+                ))
+
+        fig.update_xaxes(
+            range=[0, 7],
+            tickvals=[0, 1, 2, 3, 4, 5, 6, 7],
+        )
+
+        start_values = data['start'].sort_values()
+        end_values = data['end'].sort_values()
+
+        labels = []
+
+        for i, d in data.iterrows():
+            if i == 0:
+                labels.append(f'{hex(start_values.iloc[i])}')
+            elif i == len(data)-1:
+                labels.append(f'{hex(end_values.iloc[i])}')
+            else:
+                labels.append(f'{hex(start_values.iloc[i])}<br>{hex(end_values.iloc[i-1])}')
+
+        tickpointers = [i for i in range(len(data))]
+
+        fig.update_yaxes(
+            # tickvals=[i for i in range(len(data)+1)], 
+            tickvals = tickpointers,
+            #ticktext= labels, # Adds labels to the left-hand side of the graph
+            griddash="longdashdot",
+            gridwidth=0,
+            gridcolor="black",
+            showgrid=False,
+            showticklabels=False,
+            autorange='reversed',
+        )
+
+        fig.update_xaxes(
+            showgrid=False,
+            showticklabels=False,
+        )
+
+        fig.update_layout(
+            width=1200,
+            height=1200,
+            autosize=True,
+            margin=dict(l=200, r=20, t=20, b=20),
+            font=dict(
+                size=18,
             ),
-        ))
+            legend_title_text="Function/Locations",
+        )
 
-        ### Add top-left text with d['end']
-        # Overlapped to the right, to make it more readable
-        if pd.notna(d['overlapped_by']):
-            fig.add_trace(go.Scatter
-            (
-                x=[(x1-0.24+horizontal_gap)],
-                y=[y1-0.16],
-                text=hex(d['end']),
-                mode="text",
-                textposition="middle center",
-                marker=dict(
-                    color=fillcolor,
-                ),
-                showlegend=False,
-            ))
+        return fig
 
-            # Add bottom-left text with d['end']
-            fig.add_trace(go.Scatter
-            (
-                x=[(x1-0.24+horizontal_gap)],
-                y=[y0+0.14],
-                text=hex(d['start']),
-                mode="text",
-                textposition="middle center",
-                marker=dict(
-                    color=fillcolor,
-                ),
-                showlegend=False,
-            ))
-        else:
-            fig.add_trace(go.Scatter
-            (
-                x=[(x0+0.14+horizontal_gap)],
-                y=[y1-0.16],
-                text=hex(d['end']),
-                mode="text",
-                textposition="middle center",
-                marker=dict(
-                    color=fillcolor,
-                ),
-                showlegend=False,
-            ))
+    def write_output(fig, output_file):
+        fig.write_html(f'{output_file}.html')
 
-            ### Add bottom-left text with d['end']
-            fig.add_trace(go.Scatter
-            (
-                x=[(x0+0.14+horizontal_gap)],
-                y=[y0+0.14],
-                text=hex(d['start']),
-                mode="text",
-                textposition="middle center",
-                marker=dict(
-                    color=fillcolor,
-                ),
-                showlegend=False,
-            ))
-
-    fig.update_xaxes(
-        range=[0, 7],
-        tickvals=[0, 1, 2, 3, 4, 5, 6, 7],
-    )
-
-    start_values = data['start'].sort_values()
-    end_values = data['end'].sort_values()
-
-    labels = []
-
-    for i, d in data.iterrows():
-        if i == 0:
-            labels.append(f'{hex(start_values.iloc[i])}')
-        elif i == len(data)-1:
-            labels.append(f'{hex(end_values.iloc[i])}')
-        else:
-            labels.append(f'{hex(start_values.iloc[i])}<br>{hex(end_values.iloc[i-1])}')
-
-    tickpointers = [i for i in range(len(data))]
-
-    fig.update_yaxes(
-        # tickvals=[i for i in range(len(data)+1)], 
-        tickvals = tickpointers,
-        #ticktext= labels, # Adds labels to the left-hand side of the graph
-        griddash="longdashdot",
-        gridwidth=0,
-        gridcolor="black",
-        showgrid=False,
-        showticklabels=False,
-        autorange='reversed',
-    )
-
-    fig.update_xaxes(
-        showgrid=False,
-        showticklabels=False,
-    )
-
-    fig.update_layout(
-        width=1200,
-        height=1200,
-        autosize=True,
-        margin=dict(l=200, r=20, t=20, b=20),
-        font=dict(
-            size=18,
-        ),
-        legend_title_text="Function/Locations",
-    )
-
-    return fig
-
-def write_output(fig, output_file):
-    fig.write_html(f'{output_file}.html')
 
 if __name__ == '__main__':
     argparser = argparse.ArgumentParser()
     argparser.add_argument('--input', help='Input CSV file path', required=True, type=str)
-    argparser.add_argument('--output', help='Output HTML filename', required=False, type=str)
     args = argparser.parse_args()
 
-    if not args.output:
-        args.output = 'memory_drawer'
-
-    data = read_data(args.input)
-    fig = draw_diagram(data)
-    write_output(fig, args.output)
+    MemoryDrawer(args.input)