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
+# from herrewebpy.christianity import readplan_generator
-#readplan_generator.generate_readplan()
+# readplan_generator.generate_readplan()
 
-from herrewebpy.firmware_forensics import memory_drawer
+from herrewebpy.firmware_forensics.memory_drawer import MemoryDrawer
-memory_drawer.MemoryDrawer('sample_data/csv/stack_and_functions.csv')
+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 random, argparse, os, datetime
-import numpy as np
 import pandas as pd
 
 """
@@ -14,10 +13,37 @@ 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):
+class MemoryDrawer():
-    data = pd.read_csv(input_file)
 
-    def convert_to_int(value):
+    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)
@@ -26,8 +52,8 @@ def read_data(input_file):
             except ValueError:
                 return value 
 
-    data['start'] = data['start'].apply(convert_to_int)
+        data['start'] = data['start'].apply(_convert_to_int)
-    data['end'] = data['end'].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)
@@ -74,7 +100,7 @@ def read_data(input_file):
         return data
 
 
-def draw_diagram(data, vertical_gap_percentage=0.08, horizontal_gap=0.1):
+    def draw_diagram(data, vertical_gap_percentage=0.08, horizontal_gap=0.1):
         tickpointers = []
         labels = pd.DataFrame()
 
@@ -257,18 +283,13 @@ def draw_diagram(data, vertical_gap_percentage=0.08, horizontal_gap=0.1):
 
         return fig
 
-def write_output(fig, output_file):
+    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:
+    MemoryDrawer(args.input)
-        args.output = 'memory_drawer'
-
-    data = read_data(args.input)
-    fig = draw_diagram(data)
-    write_output(fig, args.output)