Code Editor for app.py

import streamlit as st
import numpy as np
import plotly.graph_objects as go
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.pyplot as plt
from io import BytesIO

SIZE = 256

def normalize_2d_map(map, log_scale=False):
    max = 0
    for i in range(SIZE):
        for j in range(SIZE):
            if map[i][j] > max:
                if log_scale:
                    max = log(map[i][j])
                else:
                    max = map[i][j]
    assert max != 0, "Unexpected, No data in the Map"
    for i in range(SIZE):
        for j in range(SIZE):
            map[i][j] /= max

    return map


def normalize_3d_map(map, log_scale=False):
    max = 0
    for i in range(SIZE):
        for j in range(SIZE):
            for k in range(SIZE):
                if map[i][j][k] > max:
                    if log_scale:
                        max = log(map[i][j][k])
                    else:
                        max = map[i][j][k]
    assert max != 0, "Unexpected, No data in the Map"
    for i in range(SIZE):
        for j in range(SIZE):
            for k in range(SIZE):
                map[i][j][k] /= max

    return map

def build_2d_map_pattern(f):
    map = [[0 for _ in range(SIZE)] for _ in range(SIZE)]
    try:
        init = True
        prev_byte = bytes(0)
        while True:
            curr_byte = f.read(1)
            if not curr_byte:  # EOF
                break
            if init:
                prev_byte = curr_byte
                init = False
                continue
            map[int.from_bytes(prev_byte)][int.from_bytes(curr_byte)] += 1
            prev_byte = curr_byte

        return map

    except IOError as e:
        print(f"[ERROR] While reading file: {e}")
        exit(1)


def build_3d_map_pattern(f):
    map_3d = [[[0 for _ in range(SIZE)] for _ in range(SIZE)] for _ in range(SIZE)]
    try:
        init = 2
        z_byte = y_byte = x_byte = 0
        while True:
            x_byte = f.read(1)
            if not x_byte:  # EOF
                break
            x_byte = int.from_bytes(x_byte, byteorder='big')
            if init == 2:
                y_byte = x_byte
                init -= 1
            elif init == 1:
                z_byte = y_byte
                y_byte = x_byte
                init -= 1
            else:
                map_3d[x_byte][y_byte][z_byte] += 1
                z_byte = y_byte
                y_byte = x_byte

        return map_3d

    except IOError as e:
        print(f"[ERROR] While reading file: {e}")
        exit(1)

def plot_2d_map(map_data, log_scale=False):
    if log_scale:
        map_data = np.log1p(map_data)
    
    fig, ax = plt.subplots(figsize=(10, 8))
    
    # Create a custom colormap
    colors = [(1, 1, 1, 0), (0.8, 0.8, 0.8, 1), (0, 0, 0, 1)]
    n_bins = 100
    cmap = LinearSegmentedColormap.from_list('custom_cmap', colors, N=n_bins)
    
    plt.imshow(map_data, cmap=cmap, interpolation='none')
    plt.colorbar(label='Normalized Frequency')
    plt.title('2D Representation of File Pattern')
    plt.xlabel('Current Byte')
    plt.ylabel('Previous Byte')
    
    # Convert plot to image
    buf = BytesIO()
    plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
    plt.close()
    return buf

def plot_3d_map(map_data, log_scale=False):
    if log_scale:
        map_data = np.log1p(map_data)
    
    x, y, z = [], [], []
    colors = []
    
    # Calculate total points and sample if too many
    total_points = sum(map_data[i][j][k] > 0 
                      for i in range(SIZE) 
                      for j in range(SIZE) 
                      for k in range(SIZE))
    
    MAX_POINTS = 10000  # Maximum points to display
    sample_rate = 1.0 if total_points <= MAX_POINTS else MAX_POINTS / total_points
    
    # Use numpy's random choice for sampling
    for i in range(SIZE):
        for j in range(SIZE):
            for k in range(SIZE):
                if map_data[i][j][k] > 0:
                    if np.random.random() < sample_rate:
                        x.append(i)
                        y.append(j)
                        z.append(k)
                        colors.append(map_data[i][j][k])
    
    fig = go.Figure(data=[go.Scatter3d(
        x=x, y=y, z=z,
        mode='markers',
        marker=dict(
            size=2,
            color=colors,
            colorscale='Viridis',
            opacity=0.5,
            colorbar=dict(title='Normalized Frequency')
        )
    )])
    
    fig.update_layout(
        title=f'3D Representation of File Pattern (Showing {len(x):,} of {total_points:,} points)',
        scene=dict(
            xaxis_title='Current Byte',
            yaxis_title='Previous Byte',
            zaxis_title='Previous Previous Byte'
        ),
        width=800,
        height=800
    )
    
    return fig

st.title('File Pattern Analyzer')

uploaded_file = st.file_uploader("Choose a file", type=None)
log_scale = st.checkbox('Use logarithmic scale')

if uploaded_file and st.button('Compute Patterns'):
    
    # Compute 2D pattern
    print('Computing Map')
    uploaded_file.seek(0)
    map_2d = build_2d_map_pattern(uploaded_file)
    norm_map_2d = normalize_2d_map(map_2d)
    
    # Plot 2D pattern    
    print('Building 2D Figure')
    st.subheader('2D Pattern')
    img_buf = plot_2d_map(norm_map_2d, log_scale)
    st.image(img_buf)
    
    # 3D Pattern
    uploaded_file.seek(0)
    print('Computing 3D Figure')
    st.subheader('3D Pattern')
    map_3d = build_3d_map_pattern(uploaded_file)    
    norm_map_3d = normalize_3d_map(map_3d)
    fig_3d = plot_3d_map(norm_map_3d, log_scale)
    st.plotly_chart(fig_3d)
    print('Done')