HerrewebPy/herrewebpy/firmware_forensics/memory_drawer.py

# 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)