Adding draw_boot to herrewebpy

herrewebpy/firmware_forensics/memory_drawer.py

@@ -0,0 +1,207 @@
+# Using plotly
+import plotly.graph_objects as go
+import random, argparse
+import numpy as np
+import pandas as pd
+
+def read_data(df):
+    data = pd.read_csv('stack_and_functions.csv')
+
+    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 
+
+    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'], inplace=True)
+
+    # Inverse the order of the data
+    data.reset_index(drop=True, inplace=True)
+
+    data['overlap'] = False
+
+    for i, row in data.iterrows():
+        for j, row2 in data.iterrows():
+            if i == j:
+                continue
+            if row['start'] <= row2['end'] and row['end'] > row2['start']:
+                if row['end'] - row['start'] >= row2['end'] - row2['start']:
+                    continue
+                data.at[i, 'overlap'] = True
+                data.at[j, 'overlap'] = True
+                data.at[i, 'overlap_with'] = j
+
+    data['overlap_with'] = data['overlap_with'].fillna(data.index.to_series())
+    data['overlap_with'] = data['overlap_with'].astype(float)
+
+    # Send warnings if sizes are negative
+    if (data['size'] < 0).any():
+        print(f'Warning: Negative sizes detected at indices {data[data["size"] < 0].index}')
+
+
+def draw_diagram(data):
+    tickpointers = []
+    vertical_len = len(data['overlap_with'].unique())
+    vertical_gap_percentage = 0.08
+    horizontal_gap = 0.1
+    labels = pd.DataFrame()
+
+    def random_color():
+        return f'#{random.randint(0, 0xFFFFFF):06x}'
+
+    fig = go.Figure()
+
+    for i, d in data.iterrows():
+        fillcolor = random_color()
+        data.at[i, 'fillcolor'] = fillcolor
+        
+        x0=1
+        x1=4
+
+        if d['overlap'] == False:
+            y0=d['overlap_with']
+            y1=d['overlap_with']+1
+        elif d['overlap'] == True:
+            overlaps = data.loc[data['overlap_with'] == d['overlap_with']].shape[0]
+
+            # Calculate relative size of the overlap
+            overlap_sizes = data.loc[data['overlap_with'] == d['overlap_with']].iloc[1:]['size'].sum()
+
+            if d['overlap_with'] == i:
+                y0=i
+                y1=overlaps+i
+                if i != data.shape[0]+1:
+                    if d['end'] > data.iloc[i+1].start and d['end'] < data.iloc[i+1].end:
+                        y1=overlaps+i-0.5
+                x0=x0-horizontal_gap
+                x1=x1+horizontal_gap
+            else:
+                y0=0.02+i
+                y1=0.87+i
+        else:
+            print(f'Something went wrong with {d}. Skipping')
+            continue
+
+        fig.add_shape(
+            type="rect",
+            x0=x0,
+            x1=x1,
+            y0=y0+vertical_gap_percentage,
+            y1=y1-vertical_gap_percentage,
+            line=dict(width=2),
+            fillcolor=fillcolor,
+            opacity=0.5,
+            layer="below",
+        )
+
+        # Add middle text
+        fig.add_trace(go.Scatter
+        (
+            x=[(x0+x1)/2],
+            y=[y0+0.5],
+            text=d['name'],
+            mode="text",
+            textposition="middle center",
+            name=d['name'],
+            marker=dict(
+                color=fillcolor,
+            ),
+        ))
+
+        # Add top-left text with d['end']
+        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, 5],
+        tickvals=[0, 1, 2, 3, 4, 5],
+    )
+
+    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,
+        griddash="longdashdot",
+        gridwidth=0,
+        gridcolor="black",
+        showgrid=False,
+        showticklabels=True,
+        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 being the name of the function
+        legend_title_text="Function/Locations",
+    )
+
+def write_output(fig):
+    fig.write_html("../_static/stack_and_functions.html")
+
+if __name__ == '__main__':
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument('input', help='Input CSV file path', required=True)
+    argparser.add_argument('output', help='Output HTML filename', required=False)
+    args = argparser.parse_args()
+
+    data = read_data('stack_and_functions.csv')
+    fig = draw_diagram(data)
+    write_output(fig)