HerrewebPy/herrewebpy/firmware_forensics/memory_drawer.py

# Using plotly
import plotly.graph_objects as go
import random, argparse, os, datetime
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.
"""

class MemoryDrawer():

    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:
                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')

        fig = MemoryDrawer.draw_diagram(data)
        MemoryDrawer.write_output(fig, output)


    def read_data(data):

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

        # Sort start on ascending order, and size on descending order
        data.sort_values(by=['start', 'size'], inplace=True, ascending=[True, False])

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

        data["n_overlaps"] = data["overlapped_by"].apply(lambda x: len(str(x).split(',')) if pd.notna(x) else pd.NA)
        
        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',
        )

        extra_y = 0

        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 # +(data['n_overlaps'].max() - data['n_overlaps'].min())

            if d['overlap'] == True:
                # Count number of overlaps using overlapped_by
                n_overlaps = str(d['overlapped_by']).split(',')

                # 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*d['n_overlaps'])
                    x1 = x1 - (horizontal_gap*d['n_overlaps'])

                if d['partial_overlap'] == "Bottom":
                    if pd.notna(d['partial_overlapped_by']):
                        y0 = y0 + 0.25 + (0.6**len(d['partial_overlapped_by'].split(',')))

                if d['partial_overlap'] == "Top":
                    if pd.notna(d['partial_overlapped_by']):
                        y1 = y1 - (0.6**len(d['partial_overlapped_by'].split(',')))

            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=[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']
            # 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 = tickpointers,
            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)
    args = argparser.parse_args()

    MemoryDrawer(args.input)