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

"""
This script reads a CSV file with the following columns: start,end,name,order,comment,X0,LR

- name: Name of the function or location (Required)
- start: Start address of the function or location (Required)
- end: End address of the function or location (Required)

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)

    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', '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


def draw_diagram(data, vertical_gap_percentage=0.08, horizontal_gap=0.1):
    tickpointers = []
    labels = pd.DataFrame()

    def random_color():
        return f'#{random.randint(0, 0xFFFFFF):06x}'

    fig = go.Figure()
    fig.update_layout(font=dict(family="Courier New, monospace"))

    fig.update_layout(
        plot_bgcolor='#FFFFFF',
    )

    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",
    )

    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)