Applied some minor improvements

Author: dhondta

src/exeplot/__conf__.py

@@ -1,15 +1,22 @@
 # -*- coding: UTF-8 -*-
 import logging
-import matplotlib.pyplot as plt
 import numpy
 from functools import wraps
+from warnings import filterwarnings
+
+filterwarnings("ignore", "Unable to import Axes3D.")
+
+import matplotlib.pyplot as plt
+
+
+__all__ = ["check_imports", "config", "logger", "save_figure"]
 
 
 logger = logging.getLogger("exeplot")
 config = {
     'bbox_inches':    "tight",
-#    'colormap_main':  "RdYlGn_r",
-#    'colormap_other': "jet",
+    'colormap_main':  "RdYlGn_r",
+    'colormap_other': "jet",
     'dpi':            300,
     'font_family':    "serif",
     'font_size':      10,
@@ -99,6 +106,7 @@ def _wrapper(*a, **kw):
                 ns.update(locals())
                 interact(local=ns)
             logger.info(f"Saving to {img}...")
+            plt.set_cmap("gray" if kw.get('grayscale', False) else config['colormap_main'])
             plt.savefig(img, **kw_plot)
             logger.debug(f"> saved to {img}...")
             r.append(img)

src/exeplot/__main__.py

@@ -18,6 +18,7 @@ def _setup(parser):  # pragma: no cover
     if hasattr(args, "verbose"):
         import logging
         logging.basicConfig(level=[logging.INFO, logging.DEBUG][args.verbose])
+        logging.getLogger("exeplot").level = [logging.INFO, logging.DEBUG][args.verbose]
     return args
 
 

src/exeplot/plots/__common__.py

@@ -41,7 +41,7 @@
     'cstring':  "navy",           # string table
     'const':    "cornflowerblue", # read-only data
     'literal4': "blue",           # 4-byte literal values
-    'literal4': "mediumblue",     # 8-byte literal values
+    'literal8': "mediumblue",     # 8-byte literal values
     'common':   "royalblue",      # uninitialized imported symbol definitions
 }
 MIN_ZONE_WIDTH = 3  # minimum number of samples on the entropy plot for a section (so that it can still be visible even

src/exeplot/plots/byte.py

@@ -1,32 +1,37 @@
 # -*- coding: UTF-8 -*-
 from .__common__ import Binary, COLORS
-from ..__conf__ import save_figure
+from ..__conf__ import *
 from ..utils import human_readable_size
 
 
 def arguments(parser):
     parser.add_argument("executable", help="executable sample to be plotted")
+    parser.add_argument("-g", "--grayscale", action="store_true", help="use grayscale colors instead of RGB")
     return parser
 
 
 @save_figure
-def plot(executable, height=600, **kwargs):
+def plot(executable, height=600, grayscale=False, **kwargs):
     """ draw a byte plot of the input binary """
     import matplotlib.colors as mcol
     import matplotlib.pyplot as plt
     from math import ceil, sqrt
     from matplotlib import font_manager, rcParams
     from PIL import Image, ImageDraw, ImageFont
+    # ------------------------------------------------- DRAW THE PLOT --------------------------------------------------
     # determine base variables and some helper functions
-    images, binary = [], Binary(executable)
-    n_pixels = ceil(binary.size / 3)
+    images, binary, factor = [], Binary(executable), [3, 1][grayscale]
+    n_pixels = ceil(binary.size / factor)
     s = int(ceil(sqrt(n_pixels)))
     sf = height / s
     _rgb = lambda c: tuple(map(lambda x: int(255 * x), mcol.to_rgba(c)))
+    _gs = lambda c: int(sum(a * b for a, b in zip([.299, .587, .114], _rgb(c))))
     # draw a byte plot
-    rawbytes = binary.rawbytes + int((s * s * 3) - len(binary.rawbytes)) * b'\xff'
-    images.append(Image.frombuffer("RGB", (s, s), rawbytes, "raw", "RGB", 0, 1) \
+    logger.debug("> creating the main plot from raw bytes")
+    rawbytes = binary.rawbytes + int((s * s * factor) - len(binary.rawbytes)) * b'\xff'
+    images.append(Image.frombuffer(m := ["RGB", "L"][grayscale], (s, s), rawbytes, "raw", m, 0, 1) \
                        .resize((int(s * sf), height), resample=Image.Resampling.BOX))
+    logger.debug("> creating the side plot with sections")
     if len(binary.sections) > 0:
         # map matplotlib font to PIL ImageFont path
         font_name = rcParams[f"font.{kwargs['config']['font_family']}"][0]
@@ -35,7 +40,6 @@ def plot(executable, height=600, **kwargs):
         txt_spacing = txt_h // 2
         n_lab_per_col = int((height - txt_spacing) / (txt_h + txt_spacing))
         n_cols = ceil((len(binary.sections) + 2) / n_lab_per_col)
-        #n_cols = 1
         max_txt_w = [0] * n_cols
         sections = ["Headers"] + [s for s in binary.sections] + ["Overlay"]
         draw = ImageDraw.Draw(images[0])
@@ -49,44 +53,48 @@ def plot(executable, height=600, **kwargs):
                 pass
         max_w = sum(max_txt_w) + (n_cols - 1) * txt_spacing
         # draw a separator
-        images.append(Image.new("RGB", (int(.05 * height), height), "white"))
+        images.append(Image.new(m, (int(.05 * height), height), "white"))
         # draw a sections plot aside
-        img = Image.new("RGB", (s, s), "white")
+        img = Image.new(m, (s, s), "white")
         # draw the legend with section names
-        legend = Image.new("RGB", (max_w, height), "white")
+        legend = Image.new(m, (max_w, height), "white")
         draw = ImageDraw.Draw(legend)
         _xy = lambda n, c: (txt_spacing + sum(max_txt_w[:c]) + len(max_txt_w[:c]) * txt_spacing, \
                             txt_spacing + (n % n_lab_per_col) * (txt_spacing + txt_h))
+        _c_func = [_rgb, _gs][grayscale]
         for i, name, start, end, color in binary:
             if start != end:
-                x0, y0 = min(max(ceil(((start / 3) % s)) - 1, 0), s - 1), \
-                         min(max(ceil(start / s / 3) - 1, 0), s - 1)
-                xN, yN = min(max(ceil(((end / 3) % s)) - 1, 0), s - 1), \
-                         min(max(ceil(end / s / 3) - 1, 0), s - 1)
+                x0, y0 = min(max(ceil(((start / factor) % s)) - 1, 0), s - 1), \
+                         min(max(ceil(start / s / factor) - 1, 0), s - 1)
+                xN, yN = min(max(ceil(((end / factor) % s)) - 1, 0), s - 1), \
+                         min(max(ceil(end / s / factor) - 1, 0), s - 1)
                 if y0 == yN:
                     xN = min(max(x0 + 1, xN), s - 1)
+                c = _c_func(color)
                 for x in range(x0, s if y0 < yN else xN):
-                    img.putpixel((x, y0), _rgb(color))
+                    img.putpixel((x, y0), c)
                 for y in range(y0 + 1, yN):
                     for x in range(0, s):
-                        img.putpixel((x, y), _rgb(color))
+                        img.putpixel((x, y), c)
                 if yN > y0:
                     for x in range(0, xN):
-                        img.putpixel((x, yN), _rgb(color))
+                        img.putpixel((x, yN), c)
             # fill the legend with the current section name
             if name.startswith("TOTAL"):
                 color = "black"
-            draw.text(_xy(i, ceil((i + 1) / n_lab_per_col) - 1), name, fill=_rgb(color), font=font)
+            draw.text(_xy(i, ceil((i + 1) / n_lab_per_col) - 1), name, fill=_c_func(color), font=font)
         images.append(img.resize((int(img.size[0] * sf * .2), height), resample=Image.Resampling.BOX))
-        images.append(Image.new("RGB", (int(.03 * height), height), "white"))  # draw another separator
+        images.append(Image.new(m, (int(.03 * height), height), "white"))  # draw another separator
         images.append(legend)
     # combine images horizontally
-    x, img = 0, Image.new("RGB", (sum(i.size[0] for i in images), height))
+    x, img = 0, Image.new(m, (sum(i.size[0] for i in images), height))
     for i in images:
         img.paste(i, (x, 0))
         x += i.size[0]
+    # ---------------------------------------------- CONFIGURE THE FIGURE ----------------------------------------------
     # plot combined PIL images
     #plt.tight_layout(pad=0)
+    logger.debug("> configuring the figure")
     if not kwargs.get('no_title', False):
         # set plot's title before displaying the image
         fsp = plt.gcf().subplotpars

src/exeplot/plots/diff.py

@@ -1,6 +1,6 @@
 # -*- coding: UTF-8 -*-
 from .__common__ import Binary, CACHE_DIR, COLORS, MIN_ZONE_WIDTH
-from ..__conf__ import save_figure
+from ..__conf__ import *
 
 
 def arguments(parser):
@@ -67,10 +67,15 @@ def plot(executable, executable2, legend1="", legend2="", **kwargs):
         cache = Memory(CACHE_DIR, verbose=0).cache
     except ImportError:
         cache = lambda x: x  # do nothing if joblib not installed
+    bin1, bin2 = Binary(executable), Binary(executable2)
+    if bin1.type != bin2.type:
+        raise ValueError(f"Inputs executables have different types ({bin1.type} != {bin2.type})")
     # inner function for caching sequence matches)
     @cache
     def byte_differences(bytes1, bytes2):
         return zip(*SequenceMatcher(a=bytes1, b=bytes2).get_opcodes())
+    # ------------------------------------------------- DRAW THE PLOT --------------------------------------------------
+    logger.debug("> computing the difference between executables' raw bytes")
     fs_ref = kwargs['config']['font_size']
     title = not kwargs.get('no_title', False)
     lloc = kwargs.get('legend_location', "lower right")
@@ -81,16 +86,13 @@ def byte_differences(bytes1, bytes2):
     fig.tight_layout(pad=2)
     objs[-1].axis("off")
     values, colors = {'delete': 0, 'replace': 1, 'equal': 2, 'insert': 3}, ["red", "gold", "lightgray", "green"]
-    bin1, bin2 = Binary(executable), Binary(executable2)
-    if bin1.type != bin2.type:
-        raise ValueError(f"Inputs executables have different types ({bin1.type} != {bin2.type})")
     if title:
         fig.suptitle(f"Byte-wise difference of {bin1.type} files: {bin1.basename} VS {bin2.basename}",
                      x=[.5, .55][legend1 is None], y=1, ha="center", va="bottom", **kwargs['title-font'])
     legend1, legend2 = legend1 or bin1.basename, legend2 or bin2.basename
-    (lg := kwargs['logger']).info("Matching binaries' byte sequences, this may take a while...")
+    logger.info("Matching binaries' byte sequences, this may take a while...")
     tags, alo, ahi, blo, bhi = byte_differences(bin1.rawbytes, bin2.rawbytes)
-    lg.debug("Plotting binaries...")
+    logger.debug("> plotting binaries")
     text_x = -0.012*max(bin1.size*(len(legend1)+3), bin2.size*(len(legend2)+3))
     for i, d in enumerate([(bin1, zip(tags, alo, ahi), legend1), (bin2, zip(tags, blo, bhi), legend2)]):
         binary, opcodes, label = d
@@ -129,6 +131,8 @@ def byte_differences(bytes1, bytes2):
         if len(sections) == 0:
             obj.text(.5, ref_point, "Could not parse sections", fontsize=fs_ref*1.6, color="red", ha="center",
                      va="center")
+    # ---------------------------------------------- CONFIGURE THE FIGURE ----------------------------------------------
+    logger.debug("> configuring the figure")
     cb = plt.colorbar(ScalarMappable(cmap=ListedColormap(colors, N=4)),
                       location='bottom', ax=objs[-1], fraction=0.3, aspect=50, ticks=[0.125, 0.375, 0.625, 0.875])
     cb.set_ticklabels(['removed', 'modified', 'untouched', 'added'])

src/exeplot/plots/entropy.py

@@ -1,6 +1,6 @@
 # -*- coding: UTF-8 -*-
 from .__common__ import mean, Binary, COLORS, MIN_ZONE_WIDTH, N_SAMPLES, SUBLABELS
-from ..__conf__ import save_figure
+from ..__conf__ import *
 from ..utils import shannon_entropy
 
 
@@ -25,7 +25,7 @@ def data(executable, n_samples=N_SAMPLES, window_size=lambda s: 2*s, **kwargs):
     :param n_samples:   number of samples of entropy required
     :param window_size: window size for computing the entropy
     """
-    binary = Binary(executable)
+    binary = executable if isinstance(executable, Binary) else Binary(executable)
     data = {'hash': binary.hash, 'name': binary.basename, 'size': binary.size, 'type': binary.type,
             'entropy': [], 'sections': []}
     # compute window-based entropy
@@ -94,6 +94,7 @@ def plot(*filenames, labels=None, sublabel=None, scale=False, target=None, **kwa
     from os import fstat
     if len(filenames) == 0:
         raise ValueError("No executable to plot")
+    # ------------------------------------------------- DRAW THE PLOT --------------------------------------------------
     lloc, title = kwargs.get('legend_location', "lower center"), not kwargs.get('no_title', False)
     lloc_side = lloc.split()[1] in ["left", "right"]
     nf, N_TOP, N_TOP2, N_BOT, N_BOT2 = len(filenames), 1.15, 1.37, -.15, -.37
@@ -103,12 +104,13 @@ def plot(*filenames, labels=None, sublabel=None, scale=False, target=None, **kwa
     (objs[0] if nf+[0, 1][title] > 1 else objs).axis("off")
     ref_size, ref_n, fs_ref = None, kwargs.get('n_samples', N_SAMPLES), kwargs['config']['font_size']
     for i, filepath in enumerate(filenames):
+        logger.debug(f"> plotting binary '{filepath}'")
+        binary = Executable(filepath)
         if scale and ref_size:
-            with open(filepath, "rb") as f:
-                size = fstat(f.fileno()).st_size
-            kwargs['n_samples'] = int(ref_n * size / ref_size)
+            binary.rawbytes  # triggers the computation of binary.__size
+            kwargs['n_samples'] = int(ref_n * binary.size / ref_size)
         obj = objs[i+[0, 1][title]] if nf+[0, 1][title] > 1 else objs
-        d = data(filepath, **kwargs)
+        d = data(binary, **kwargs)
         n, label = len(d['entropy']), None
         if not ref_size:
             ref_size = d['size']
@@ -180,13 +182,16 @@ def plot(*filenames, labels=None, sublabel=None, scale=False, target=None, **kwa
             # draw entropy
             obj.plot(x, d['entropy'][start:end+1], c=c, zorder=10, lw=.1)
             obj.fill_between(x, [0] * len(x), d['entropy'][start:end+1], facecolor=c)
-            l = obj.hlines(y=mean(d['entropy'][start:end+1]), xmin=x[0], xmax=x[-1], color="black", linewidth=.5, linestyle=(0, (5, 5)))
+            l = obj.hlines(y=mean(d['entropy'][start:end+1]), xmin=x[0], xmax=x[-1], color="black", linewidth=.5,
+                           linestyle=(0, (5, 5)))
             i += 1
         if len(d['sections']) > 0:
             l.set_label("Average entropy of section")
         else:
             obj.text(.5, ref_point, "Could not parse sections", fontsize=fs_ref*1.6, color="red", ha="center",
                      va="center")
+    # ---------------------------------------------- CONFIGURE THE FIGURE ----------------------------------------------
+    logger.debug("> configuring the figure")
     plt.subplots_adjust(left=[.15, .02][labels is None and sublabel is None], right=[1.02, .82][lloc_side],
                         bottom=.5/max(1.75, nf))
     h, l = (objs[[0, 1][title]] if nf+[0, 1][title] > 1 else objs).get_legend_handles_labels()

src/exeplot/plots/graph.py

@@ -2,16 +2,11 @@
 import matplotlib.pyplot as plt
 
 from .__common__ import Binary, CACHE_DIR, COLORS, MIN_ZONE_WIDTH
-from ..__conf__ import check_imports, save_figure
-
-try:
-    # dirty fix to known issue with angr: AttributeError: module 'unicorn' has no attribute 'UC_ARCH_RISCV'.
-    __import__("unicorn").UC_ARCH_RISCV = 8
-except ModuleNotFoundError: # pragma: no cover
-    pass  # 'unicorn' is an optional dependency of 'angr' ; fix it only if it is already installed
+from ..__conf__ import *
 
 check_imports("angr", "networkx", "pygraphviz")
 
+
 _DEFAULT_ALGORITHM, _DEFAULT_ENGINE = "fast", "default"
 _ENGINES = ["default", "pcode", "vex"]
 
@@ -32,6 +27,8 @@ def plot(executable, algorithm=_DEFAULT_ALGORITHM, engine=_DEFAULT_ENGINE, **kwa
     import networkx as nx
     import pygraphviz as pgv
     from math import ceil, log2
+    # ------------------------------------------------- DRAW THE PLOT --------------------------------------------------
+    logger.debug("> computing the executable's CFG")
     engine = {k: getattr(angr.engines, "UberEngine" if k != "pcode" else f"UberEngine{k.capitalize()}") \
               for k in _ENGINES}[engine]
     project = angr.Project(executable, auto_load_libs=False, engine=engine)
@@ -40,6 +37,8 @@ def plot(executable, algorithm=_DEFAULT_ALGORITHM, engine=_DEFAULT_ENGINE, **kwa
     for node in cfg.graph.nodes():
         labels[node] = f"{node.name}\n0x{node.addr:x}" if hasattr(node, "name") and node.name else f"0x{node.addr:x}"
         node_colors.append("red" if node.function_address == node.addr else "lightblue")
+    # ---------------------------------------------- CONFIGURE THE FIGURE ----------------------------------------------
+    logger.debug("> configuring the figure")
     n = max(10, min(30, ceil(log2(n_nodes := len(cfg.graph.nodes()) + 1) * 2)))
     fig = plt.figure(figsize=(n, n))
     nx.draw(cfg.graph, nx.kamada_kawai_layout(cfg.graph), fig.gca(), font_size=8, with_labels=True, labels=labels,

src/exeplot/plots/nested_pie.py

@@ -1,6 +1,6 @@
 # -*- coding: UTF-8 -*-
 from .__common__ import Binary, COLORS, SHADOW
-from ..__conf__ import save_figure
+from ..__conf__ import *
 from ..utils import human_readable_size
 
 
@@ -14,19 +14,24 @@ def plot(executable, **kwargs):
     """ draw a nested pie chart of segments (if relevant) and sections (including overlaps) of the input binary """
     import matplotlib.pyplot as plt
     from math import ceil
+    # ------------------------------------------------- DRAW THE PLOT --------------------------------------------------
     binary = Binary(executable)
     fig, ax = plt.subplots(figsize=(6, 3), subplot_kw=dict(aspect="equal"))
     pie_kw = {'shadow': SHADOW} if kwargs['config']['shadow'] else {}
     seg_layers = sum(1 for _ in binary._data(segments=True, overlap=True))
     if binary.type != "PE":
+        logger.debug("> computing pie segments")
         for i, x, w, labels, colors, legend in binary._data(segments=True, overlap=True):
             size = .4 / seg_layers
             ax.pie(w, radius=.45-i*size, colors=colors, startangle=90,
                    wedgeprops={'width': size, 'edgecolor': "w", 'linewidth': 1}, **pie_kw)
+    logger.debug("> computing pie sections")
     sec_layers = sum(1 for _ in binary._data(overlap=True))
     for i, x, w, labels, colors, legend in binary._data(overlap=True):
         size = .42 / sec_layers
         ax.pie(w, radius=1-i*size, colors=colors, startangle=90, wedgeprops={'width': size}, **pie_kw)
+    # ---------------------------------------------- CONFIGURE THE FIGURE ----------------------------------------------
+    logger.debug("> configuring the figure")
     ncols = ceil(len(legend['colors']) / 12)
     ax.legend([plt.Rectangle((0, 0), 1, 1, color=c) for c in legend['colors']], legend['texts'], loc="center left",
               bbox_to_anchor=(1, 0, 0.5, 1), ncol=ncols, fontsize=ceil(kwargs['config']['font_size']*.7))

src/exeplot/plots/pie.py

@@ -1,6 +1,6 @@
 # -*- coding: UTF-8 -*-
 from .__common__ import Binary, COLORS, SHADOW
-from ..__conf__ import save_figure
+from ..__conf__ import *
 from ..utils import human_readable_size
 
 
@@ -15,6 +15,8 @@ def plot(executable, donut=False, **kwargs):
     """ draw a pie chart of the sections of the input binary """
     import matplotlib.pyplot as plt
     from math import ceil
+    # ------------------------------------------------- DRAW THE PLOT --------------------------------------------------
+    logger.debug("> computing pie sections")
     fs_ref = kwargs['config']['font_size']
     binary = Binary(executable)
     fig, ax = plt.subplots(figsize=(6, 3), subplot_kw=dict(aspect="equal"))
@@ -27,6 +29,8 @@ def plot(executable, donut=False, **kwargs):
     pie_kw = {'shadow': SHADOW} if kwargs['config']['shadow'] else {}
     _, texts = ax.pie(data, colors=colors, labels=labels, textprops=txt_kw, labeldistance=.55, startangle=90,
                            wedgeprops={'width': [1., .55][donut]}, **pie_kw)
+    # ---------------------------------------------- CONFIGURE THE FIGURE ----------------------------------------------
+    logger.debug("> configuring the figure")
     ax.legend([plt.Rectangle((0, 0), 1, 1, color=c) for c in legend['colors']], legend['texts'], loc="center left",
               bbox_to_anchor=(1, 0, 0.5, 1), ncol=ncols, fontsize=ceil(fs_ref*.7))
     plt.setp(texts, size=fs_ref*.8, weight="bold")