refactor: yolodfu detection, pongo compression

Update run.c

src/m8/checkm8.c

@@ -10,6 +10,7 @@
 
 #include "../utils.h"
 #include "../usb/shim.h"
+#include "../usb/pongo_helper.h"
 
 static struct {
     uint8_t b_len, b_descriptor_type;
@@ -162,12 +163,8 @@ bool checkm8_stage_patch(const usb_handle_t *handle, struct DeviceConfiguration
     uint64_t* p = (uint64_t*)blank;
     p[5] = payloadConfig->insecure_memory_base;
 
-    create_pongo_payload_for_device(deviceConfig->cpid, &payload, &data_sz);
-
-    if (payload == NULL || data_sz == 0) {
-        LOG_ERROR("unsupported cpid 0x%X", (unsigned)deviceConfig->cpid);
-        return false;
-    }
+    ret = create_pongo_payload_for_device(deviceConfig->cpid, &payload, &data_sz);
+    if (!ret) return false;
 
     LOG_VERBOSE("setting up stage 2 for CPID 0x%X", (unsigned)deviceConfig->cpid);
 
@@ -200,3 +197,44 @@ bool checkm8_stage_patch(const usb_handle_t *handle, struct DeviceConfiguration
     free(data);
     return ret;
 }
+
+bool checkm8_boot_pongo(usb_handle_t *handle)
+{
+    transfer_ret_t transfer_ret;
+
+    uint8_t *out;
+    size_t out_len;
+
+    if (!prepare_pongo(&out, &out_len)) {
+        LOG_ERROR("Failed to prepare pongoOS payload");
+        return false;
+    }
+
+    {
+        size_t len = 0;
+
+        while (len < out_len)
+        {
+            size_t chunk = (out_len - len > 0x800) ? 0x800 : (out_len - len);
+
+        retry:
+            send_usb_control_request(handle, 0x21, DFU_DNLOAD, 0, 0, out + len, chunk, &transfer_ret);
+
+            if (transfer_ret.sz != chunk || transfer_ret.ret != USB_TRANSFER_OK)
+            {
+                LOG_VERBOSE("retrying at len = %zu", len);
+                sleep_ms(100);
+                goto retry;
+            }
+
+            len += chunk;
+            LOG_VERBOSE("len = %zu", len);
+        }
+    }
+
+    free(out);
+
+    send_usb_control_request_no_data(handle, 0x21, 4, 0, 0, 0, NULL);
+    LOG_SUCCESS("PongoOS has been sent off");
+    return true;
+}

src/m8/checkm8.h

@@ -18,5 +18,6 @@ bool checkm8_stage_reset(const usb_handle_t *handle);
 bool checkm8_stage_setup(const usb_handle_t *handle, struct DeviceConfiguration *deviceConfig);
 bool checkm8_stage_spray(const usb_handle_t *handle, struct DeviceConfiguration *deviceConfig);
 bool checkm8_stage_patch(const usb_handle_t *handle, struct DeviceConfiguration *deviceConfig, struct PayloadConfiguration *payloadConfig);
+bool checkm8_boot_pongo(usb_handle_t *handle);
 
 #endif // CHECKM8_H

src/m8/payload.c

@@ -237,12 +237,12 @@ static const yolo_payload_t yolo_payloads[] = {
     {0x7000, payloads_yolo_t7000_bin, payloads_yolo_t7000_bin_len},
 };
 
-void create_pongo_payload_for_device(
+bool create_pongo_payload_for_device(
     uint16_t cpid,
     const uint8_t **payload,
     size_t *payloadSize)
 {
-    LOG_VERBOSE("Preparing YoloDFU payload for CPID 0x%x.", cpid);
+    LOG_VERBOSE("Preparing YoloDFU payload for CPID 0x%x", cpid);
 
     *payload = NULL;
     *payloadSize = 0;
@@ -253,9 +253,10 @@ void create_pongo_payload_for_device(
         if (yolo_payloads[i].cpid == cpid) {
             *payload = yolo_payloads[i].data;
             *payloadSize = yolo_payloads[i].len;
-            return;
+            return true;
         }
     }
 
-    LOG("Failed to prepare payload for device with CPID 0x%x.", cpid);
+    LOG("Failed to prepare payload for device with CPID 0x%x", cpid);
+    return false;
 }

src/m8/payload.h

@@ -46,6 +46,6 @@ typedef struct {
 
 bool checkm8_find_device_configuration_for_cpid(uint16_t cpid, struct DeviceConfiguration *config);
 bool checkm8_find_payload_configuration_for_cpid(uint16_t cpid, struct PayloadConfiguration *config);
-void create_pongo_payload_for_device(uint16_t cpid, const uint8_t **payload, size_t *payloadSize);
+bool create_pongo_payload_for_device(uint16_t cpid, const uint8_t **payload, size_t *payloadSize);
 
 #endif // PAYLOAD_H

src/m8/run.c

@@ -19,164 +19,161 @@
 #include "dfu.h"
 #include "payload.h"
 
+#define VID_APPLE   (0x5AC)
+#define PID_PONGO   (0x4141)
+#define PID_DFU     (0x1227)
+
 static void *exploit_thread(void *arg)
 {
+    LOG_VERBOSE("Starting thread for DFU devices");
     LOG("Waiting for DFU devices...");
+
     shared_t *s = (shared_t *)arg;
 
     usb_handle_t handle;
     struct DeviceConfiguration deviceConfig;
     struct PayloadConfiguration payloadConfig;
 
-    int stage = STAGE_PREPARE;
+    int stage = STAGE_RESET;
     bool ret = false;
 
-    init_usb_handle(&handle, 0x5AC, 0x1227);
+    init_usb_handle(&handle, VID_APPLE, PID_DFU);
 
-    while (stage != STAGE_JAILBREAK && !atomic_load(&s->stop)) {
+    while (!atomic_load(&s->stop)) {
         if (!wait_usb_handle(&handle)) {
-            atomic_store(&s->result, -1);
-            atomic_store(&s->stop, true);
-            break;
+            reset_usb_handle(&handle);
+            continue;
         }
 
+        char *sn = get_usb_serial_number(&handle);
+        if (!sn) {
+            reset_usb_handle(&handle);
+            continue;
+        }
+
+        uint64_t cpid = dfu_serial_number_get_cpid(sn);
+        if (!checkm8_find_device_configuration_for_cpid(cpid, &deviceConfig)
+            || !checkm8_find_payload_configuration_for_cpid(cpid, &payloadConfig))
+        {
+            free(sn);
+            reset_usb_handle(&handle);
+            continue;
+        }
+
+        if (dfu_serial_number_is_pwned(sn)) {
+            LOG_WARN("Detected already PWN'd device, skipping");
+            free(sn);
+            reset_usb_handle(&handle);
+            continue;
+        } else if (dfu_serial_number_is_in_dfu_mode(sn)) {
+            LOG_SUCCESS("Detected DFU device");
+        } else if (dfu_serial_number_is_in_yolo_dfu(sn)) {
+            LOG_SUCCESS("Detected YoloDFU device");
+            stage = STAGE_YOLODFU;
+        } else {
+            free(sn);
+            reset_usb_handle(&handle);
+            continue;
+        }
+
+        free(sn);
         atomic_store(&s->stage, stage);
 
         switch (stage) {
-        case STAGE_PREPARE: {
-            char *sn = get_usb_serial_number(&handle);
-            if (!sn) goto fail;
-
-            uint64_t cpid = dfu_serial_number_get_cpid(sn);
-
-            if (!checkm8_find_device_configuration_for_cpid(cpid, &deviceConfig)) {
-                free(sn);
-                goto fail;
-            }
-
-            if (!checkm8_find_payload_configuration_for_cpid(cpid, &payloadConfig)) {
-                free(sn);
-                goto fail;
-            }
-
-            if (!dfu_serial_number_is_in_dfu_mode(sn)) {
-                free(sn);
-                goto fail;
-            }
-
-            free(sn);
-            stage = STAGE_RESET;
-            break;
-        }
-        case STAGE_RESET:
-            ret = checkm8_stage_reset(&handle);
-            stage = ret ? STAGE_SETUP : STAGE_PREPARE;
-            break;
-        case STAGE_SETUP:
-            LOG("Setting up the device for exploitation");
-            ret = checkm8_stage_setup(&handle, &deviceConfig);
-            stage = ret ? STAGE_SPRAY : STAGE_PREPARE;
-            break;
-        case STAGE_SPRAY:
-            LOG("Spraying the device with USB requests");
-            ret = checkm8_stage_spray(&handle, &deviceConfig);
-            stage = ret ? STAGE_PATCH : STAGE_PREPARE;
-            break;
-        case STAGE_PATCH:
-            LOG("Right before trigger (bug setup)");
-            ret = checkm8_stage_patch(&handle, &deviceConfig, &payloadConfig);
-            if (!ret) goto fail;
-            stage = STAGE_PONGO;
-            break;
-        case STAGE_PONGO:
-            LOG("Booting pongoOS");
-
-            close_usb_handle(&handle);
-            sleep_ms(3000);
-
-            checkm8_boot_pongo(&handle);
-
-            atomic_store(&s->exploit_done, true);
-            atomic_store(&s->pongo_done, true);
-            atomic_store(&s->stage, STAGE_JAILBREAK);
-
-            stage = STAGE_JAILBREAK;
-            break;
-        default:
-            goto fail;
+            case STAGE_RESET:
+                ret = checkm8_stage_reset(&handle);
+                stage = ret ? STAGE_SETUP : STAGE_RESET;
+                break;
+            case STAGE_SETUP:
+                LOG("Checking if device is ready");
+                ret = checkm8_stage_setup(&handle, &deviceConfig);
+                stage = ret ? STAGE_SPRAY : STAGE_RESET;
+                break;
+            case STAGE_SPRAY:
+                LOG("Setting up the exploit (this is the heap spray)");
+                ret = checkm8_stage_spray(&handle, &deviceConfig);
+                stage = ret ? STAGE_PATCH : STAGE_RESET;
+                break;
+            case STAGE_PATCH:
+                LOG("Right before trigger (this is the real bug setup)");
+                ret = checkm8_stage_patch(&handle, &deviceConfig, &payloadConfig);
+                stage = STAGE_RESET;
+                if (!ret) {
+                    LOG_ERROR("Failed doing trigger");
+                    break;
+                }
+                continue;
+            case STAGE_YOLODFU:
+                LOG("Booting PongoOS...");
+                ret = checkm8_boot_pongo(&handle);
+                if (!ret) LOG_ERROR("Failed to send PongoOS to device");
+                stage = STAGE_RESET;
+                break;
         }
 
         reset_usb_handle(&handle);
     }
 
-    atomic_store(&s->result, 1);
-    return NULL;
-
-fail:
-    atomic_store(&s->result, -1);
-    atomic_store(&s->stop, true);
+    close_usb_handle(&handle);
     return NULL;
 }
 
 static void *pongo_thread(void *arg)
 {
-    LOG("Waiting for Pongo devices...");
+    LOG_VERBOSE("Starting thread for PongoOS devices");
+
     shared_t *s = (shared_t *)arg;
 
     usb_handle_t handle;
     bool seen = false;
 
-    init_usb_handle(&handle, 0x5AC, 0x4141);
+    init_usb_handle(&handle, VID_APPLE, PID_PONGO);
 
     while (!atomic_load(&s->stop)) {
-
         if (!wait_usb_handle(&handle)) {
-            atomic_store(&s->result, -1);
-            atomic_store(&s->stop, true);
-            break;
+            if (atomic_load(&s->stop)) break;
+            reset_usb_handle(&handle);
+            continue;
         }
 
-        if (!seen) {
-            seen = true;
+        if (seen) continue;
+        seen = true;
 
-            LOG_SUCCESS("Pongo device detected!");
+        LOG_SUCCESS("Detected PongoOS device");
 
-            if (palerain_flags & palerain_option_pongo_exit) {
-                close_usb_handle(&handle);
-                atomic_store(&s->pongo_done, true);
-                atomic_store(&s->stage, STAGE_DONE);
-                break;
-            }
+        atomic_store(&s->stage, STAGE_PONGO);
 
-            char paleinfo[64];
-            snprintf(paleinfo, sizeof(paleinfo), "palera1n_flags 0x%" PRIx64, palerain_flags);
-
-            char xargs_cmd[0x270];
-            snprintf(xargs_cmd, sizeof(xargs_cmd), "xargs %s", boot_args);
-            if ((palerain_flags & palerain_option_setup_rootful)) {
-                strncat(xargs_cmd, " wdt=-1", 0x270 - strlen(xargs_cmd) - 1);
-            }
-
-            LOG("Boot arguments: %s", xargs_cmd);
-
-            issue_pongo_command(&handle, "fuse lock");
-            issue_pongo_command(&handle, "sep auto");
-            upload_buffer_to_pongo(&handle, g_payload_kpf.data, g_payload_kpf.data_len);
-            issue_pongo_command(&handle, "modload");
-            issue_pongo_command(&handle, paleinfo);
-            upload_buffer_to_pongo(&handle, g_payload_ramdisk.data, g_payload_ramdisk.data_len);
-            issue_pongo_command(&handle, "ramdisk");
-            upload_buffer_to_pongo(&handle, g_payload_overlay.data, g_payload_overlay.data_len);
-            issue_pongo_command(&handle, "overlay");
-            if (strlen(boot_args) > 0)
-                issue_pongo_command(&handle, xargs_cmd);
-            issue_pongo_command(&handle, "bootx");
-
-            atomic_store(&s->pongo_done, true);
+        if (palerain_flags & palerain_option_pongo_exit) {
+            close_usb_handle(&handle);
             atomic_store(&s->stage, STAGE_DONE);
-
-            break;
+            atomic_store(&s->stop, true);
+            return NULL;
         }
+
+        char paleinfo[64];
+        snprintf(paleinfo, sizeof(paleinfo), "palera1n_flags 0x%" PRIx64, palerain_flags);
+
+        char xargs_cmd[0x270];
+        snprintf(xargs_cmd, sizeof(xargs_cmd), "xargs %s", boot_args);
+        if (palerain_flags & palerain_option_setup_rootful) {
+            strncat(xargs_cmd, " wdt=-1", sizeof(xargs_cmd) - strlen(xargs_cmd) - 1);
+        }
+
+        issue_pongo_command(&handle, "fuse lock");
+        issue_pongo_command(&handle, "sep auto");
+        upload_buffer_to_pongo(&handle, g_payload_kpf.data, g_payload_kpf.data_len);
+        issue_pongo_command(&handle, "modload");
+        issue_pongo_command(&handle, paleinfo);
+        upload_buffer_to_pongo(&handle, g_payload_ramdisk.data, g_payload_ramdisk.data_len);
+        issue_pongo_command(&handle, "ramdisk");
+        upload_buffer_to_pongo(&handle, g_payload_overlay.data, g_payload_overlay.data_len);
+        issue_pongo_command(&handle, "overlay");
+        if (strlen(boot_args) > 0) issue_pongo_command(&handle, xargs_cmd);
+        issue_pongo_command(&handle, "bootx");
+
+        atomic_store(&s->stage, STAGE_DONE);
+        atomic_store(&s->stop, true);
+        break;
     }
 
     close_usb_handle(&handle);
@@ -193,7 +190,5 @@ bool exploit(shared_t *state)
     pthread_join(t1, NULL);
     pthread_join(t2, NULL);
 
-    return atomic_load(&state->result) == 1 &&
-           atomic_load(&state->exploit_done) &&
-           atomic_load(&state->pongo_done);
+    return atomic_load(&state->result) == 1 && atomic_load(&state->stop);
 }

src/m8/run.h

@@ -13,23 +13,18 @@ extern "C" {
 #endif
 
 enum exploit_stage {
-    STAGE_PREPARE = 0,
-    STAGE_RESET = 1,
-    STAGE_SETUP = 2,
-    STAGE_SPRAY = 3,
-    STAGE_PATCH = 4,
-    STAGE_PONGO = 5,
-    STAGE_JAILBREAK = 6,
-    STAGE_DONE = 7,
+    STAGE_RESET     = 0,
+    STAGE_SETUP     = 1,
+    STAGE_SPRAY     = 2,
+    STAGE_PATCH     = 3,
+    STAGE_YOLODFU   = 4,
+    STAGE_PONGO     = 5,
+    STAGE_DONE      = 6,
 };
 
 typedef struct {
     atomic_int result;
     atomic_bool stop;
-
-    atomic_bool exploit_done;
-    atomic_bool pongo_done;
-
     atomic_int stage;
 } shared_t;
 

src/main.cpp

@@ -1,9 +1,10 @@
-#include <unistd.h> // isatty, STDIN_FILENO
-#include <cstdlib>  // getenv
+#include <unistd.h>
+#include <cstdlib>
 #include <iostream>
 #include <getopt.h>
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
 
 #ifdef WITH_GUI
 # include <wx/wx.h>

src/usb/pongo_helper.c

@@ -14,57 +14,72 @@
 
 #include "../gen/payloads/lz4dec.h"
 
-void compress_pongo(void *out, size_t *out_len) {
-    size_t len = g_payload_pongo.data_len;
-    size_t out_len_ = *out_len;
-    *out_len = LZ4_compress_HC(g_payload_pongo.data, out, len, out_len_, LZ4HC_CLEVEL_MAX);
+static bool compress_pongo(uint8_t **out, size_t *out_len)
+{
+    const uint8_t *src = g_payload_pongo.data;
+    int src_size = (int)g_payload_pongo.data_len;
+
+    int max_size = LZ4_compressBound(src_size);
+
+    uint8_t *buf = malloc(max_size);
+    if (!buf) return false;
+
+    int size = LZ4_compress_HC(
+        (const char *)src,
+        (char *)buf,
+        src_size,
+        max_size,
+        LZ4HC_CLEVEL_MAX
+    );
+
+    if (size <= 0) {
+        free(buf);
+        return false;
+    }
+
+    *out = buf;
+    *out_len = (size_t)size;
+
+    return true;
 }
 
-void checkm8_boot_pongo(usb_handle_t *handle) {
-    transfer_ret_t transfer_ret;
-    LOG_VERBOSE("Booting pongoOS");
-    LOG_VERBOSE("Compressing pongoOS");
-    LOG_VERBOSE("Appending shellcode to the top of pongoOS (512 bytes)");
-    void *shellcode = malloc(512);
-    memcpy(shellcode, payloads_lz4dec_bin, payloads_lz4dec_bin_len);
-    size_t out_len = g_payload_pongo.data_len;
-    void *out = malloc(out_len);
-    compress_pongo(out, &out_len);
-    LOG_VERBOSE("Compressed pongoOS from %u to %zu bytes", g_payload_pongo.data_len, out_len);
-    void *tmp = malloc(out_len + 512);
-    memcpy(tmp, shellcode, 512);
-    memcpy(tmp + 512, out, out_len);
-    free(out);
-    out = tmp;
-    out_len += 512;
-    free(shellcode);
-    LOG_VERBOSE("Setting the compressed size into the shellcode");
-    uint32_t* size = (uint32_t*)(out + 0x1fc);
-    *size = out_len - 512;
-    LOG_VERBOSE("Reconnecting to device");
-    init_usb_handle(handle, 0x5AC, 0x1227);
-    LOG_VERBOSE("Waiting for device to be ready");
-    wait_usb_handle(handle);
-    {
-        size_t len = 0;
-        size_t size;
-        while(len < out_len)
-        {
-        retry:
-            size = ((out_len - len) > 0x800) ? 0x800 : (out_len - len);
-            send_usb_control_request(handle, 0x21, DFU_DNLOAD, 0, 0, (unsigned char*)&out[len], size, &transfer_ret);
-            if(transfer_ret.sz != size || transfer_ret.ret != USB_TRANSFER_OK)
-            {
-                LOG_VERBOSE("retrying at len = %zu", len);
-                sleep_ms(100);
-                goto retry;
-            }
-            len += size;
-            LOG_VERBOSE("len = %zu", len);
-        }
+#define SHELLCODE_SZ (512)
+#define PONGO_MAX_SZ (0x7fe00)
+
+bool prepare_pongo(uint8_t **out, size_t *out_len)
+{
+    if (g_payload_pongo.data_len > PONGO_MAX_SZ) {
+        LOG_ERROR("Pongo payload too large: %zu bytes (max 511 KB)", g_payload_pongo.data_len);
+        return false;
     }
-    send_usb_control_request_no_data(handle, 0x21, 4, 0, 0, 0, NULL);
-    LOG_VERBOSE("pongoOS sent, should be booting");
+    
+    size_t pongoSize;
+    uint8_t *pongo;
+
+    if (!compress_pongo(&pongo, &pongoSize))
+        return false;
+
+    uint8_t *payload = malloc(SHELLCODE_SZ + pongoSize);
+    if (!payload) {
+        free(pongo);
+        return false;
+    }
+
+    memcpy(payload, payloads_lz4dec_bin, SHELLCODE_SZ);
+    memcpy(payload + SHELLCODE_SZ, pongo, pongoSize);
+
+    free(pongo);
+
+    pongo = payload;
+    pongoSize += SHELLCODE_SZ;
+
+    uint32_t *pongoSizeInData = (uint32_t *)(pongo + 0x1fc);
+    *pongoSizeInData = (uint32_t)pongoSize - SHELLCODE_SZ;
+
+    *out = pongo;
+    *out_len = pongoSize;
+
+    return true;
 }
 
 int issue_pongo_command(const usb_handle_t *handle, const char *command) {

src/usb/pongo_helper.h

@@ -5,8 +5,7 @@
 
 #define CMD_LENGTH_MAX 512
 
-void compress_pongo(void *out, size_t *out_len);
-void checkm8_boot_pongo(usb_handle_t *handle);
+bool prepare_pongo(uint8_t **out, size_t *out_len);
 int issue_pongo_command(const usb_handle_t *handle, const char *command);
 bool upload_buffer_to_pongo(usb_handle_t *handle, const void *data, size_t length);
 

src/wx/DfuPanel.cpp

@@ -88,10 +88,9 @@ DfuPanel::DfuPanel(MainFrame* frame, wxWindow* parent)
     });
 
     auto* bottomRow = new wxBoxSizer(wxHORIZONTAL);
-
     bottomRow->AddStretchSpacer();
-    bottomRow->Add(m_backButton, 0, wxALL, 12);
-    bottomRow->Add(m_startButton, 0, wxALL, 12);
+    bottomRow->Add(m_backButton, 0, wxRIGHT | wxLEFT | wxBOTTOM, 12);
+    bottomRow->Add(m_startButton, 0, wxRIGHT | wxLEFT | wxBOTTOM, 12);
 
     root->Add(bottomRow, 0, wxEXPAND);
 

src/wx/ExploitPanel.cpp

@@ -9,34 +9,31 @@
 
 #include "AppFrame.hpp"
 #include "../m8/run.h"
+#include "../utils.h"
 
 static const char* stage_to_string(int stage)
 {
     switch (stage) {
-        case STAGE_PREPARE: return "Checking if device is ready";
-        case STAGE_RESET:   return "Checking if device is ready";
-        case STAGE_SETUP:   return "Setting up the exploit (this is the heap spray)";
-        case STAGE_SPRAY:   return "Setting up the exploit (this is the heap spray)";
-        case STAGE_PATCH:   return "Right before trigger (this is the real bug setup)";
-        case STAGE_PONGO:   return "Booting...";
-        case STAGE_JAILBREAK: return "Booting...";
-        case STAGE_DONE:    return "Done.";
-        default:            return "Working...";
+        case STAGE_SETUP:       return "Checking if device is ready";
+        case STAGE_SPRAY:       return "Setting up the exploit (this is the heap spray)";
+        case STAGE_PATCH:       return "Right before trigger (this is the real bug setup)";
+        case STAGE_YOLODFU:     return "Booting PongoOS...";
+        case STAGE_PONGO:       return "Booting...";
+        case STAGE_DONE:        return "Done.";
+        default:                return "Idle.";
     }
 }
 
 static int stage_to_progress(int stage)
 {
     switch (stage) {
-        case STAGE_PREPARE: return 5;
-        case STAGE_RESET:   return 15;
-        case STAGE_SETUP:   return 30;
-        case STAGE_SPRAY:   return 50;
-        case STAGE_PATCH:   return 70;
-        case STAGE_PONGO:   return 90;
-        case STAGE_JAILBREAK: return 95;
-        case STAGE_DONE:    return 100;
-        default:            return 0;
+        case STAGE_SETUP:       return 30;
+        case STAGE_SPRAY:       return 40;
+        case STAGE_PATCH:       return 50;
+        case STAGE_YOLODFU:     return 60;
+        case STAGE_PONGO:       return 70;
+        case STAGE_DONE:        return 100;
+        default:                return 0;
     }
 }
 
@@ -67,7 +64,11 @@ ExploitPanel::ExploitPanel(MainFrame* frame, wxWindow* parent)
     root->Add(m_progressBar, 0, wxALL | wxEXPAND, 20);
 
     root->AddStretchSpacer();
-    root->Add(m_backButton, 0, wxALL | wxALIGN_RIGHT, 12);
+    auto* bottomRow = new wxBoxSizer(wxHORIZONTAL);
+    bottomRow->AddStretchSpacer();
+    bottomRow->Add(m_backButton, 0, wxRIGHT | wxLEFT | wxBOTTOM, 12);
+
+    root->Add(bottomRow, 0, wxEXPAND);
 
     SetSizer(root);
 }
@@ -82,7 +83,6 @@ void ExploitPanel::SetDeviceState(const DeviceState& state)
 
 void ExploitPanel::StartExploit()
 {
-    m_statusText->SetLabel("Starting...");
     m_backButton->Disable();
 
     std::thread([this]()
@@ -91,38 +91,36 @@ void ExploitPanel::StartExploit()
         if (!deviceState.connected || deviceState.mode != DeviceMode::DFU)
             return;
 
-        shared_t state{};
-        atomic_init(&state.result, 0);
-        atomic_init(&state.stop, false);
-        atomic_init(&state.exploit_done, false);
-        atomic_init(&state.pongo_done, false);
-        atomic_init(&state.stage, STAGE_PREPARE);
+        auto state = std::make_shared<shared_t>();
+        atomic_init(&state->stage, STAGE_RESET);
+        atomic_init(&state->stop, false);
+        atomic_init(&state->result, 0);
 
-        std::thread worker([&]() {
-            exploit(&state);
+        std::thread worker([state]()
+        {
+            exploit(state.get());
         });
 
-        worker.detach();
-
         while (true)
         {
-            int stage = atomic_load(&state.stage);
+            int stage = atomic_load(&state->stage);
             int progress = stage_to_progress(stage);
 
-            wxTheApp->CallAfter([this, stage, progress]() {
+            wxTheApp->CallAfter([this, stage, progress]()
+            {
                 m_statusText->SetLabel(stage_to_string(stage));
                 m_progressBar->SetValue(progress);
+                if (stage == STAGE_DONE) m_backButton->Enable();
             });
 
-            if (stage == STAGE_DONE) {
-                wxTheApp->CallAfter([this]() {
-                    m_backButton->Enable();
-                });
+            if (stage == STAGE_DONE)
                 break;
-            }
 
             std::this_thread::sleep_for(std::chrono::milliseconds(100));
         }
+
+        if (worker.joinable())
+            worker.join();
     }).detach();
 }
 

src/wx/RecoveryPanel.cpp

@@ -61,8 +61,8 @@ RecoveryPanel::RecoveryPanel(MainFrame* frame, wxWindow* parent)
 
     auto* bottomRow = new wxBoxSizer(wxHORIZONTAL);
     bottomRow->AddStretchSpacer();
-    bottomRow->Add(m_backButton, 0, wxALL, 12);
-    bottomRow->Add(m_nextButton, 0, wxALL, 12);
+    bottomRow->Add(m_backButton, 0, wxRIGHT | wxLEFT | wxBOTTOM, 12);
+    bottomRow->Add(m_nextButton, 0, wxRIGHT | wxLEFT | wxBOTTOM, 12);
 
     root->AddStretchSpacer();
     root->Add(bottomRow, 0, wxEXPAND);

src/wx/SettingsPanel.cpp

@@ -95,7 +95,7 @@ SettingsPanel::SettingsPanel(MainFrame* frame, wxWindow* parent)
 
     auto* bottomRow = new wxBoxSizer(wxHORIZONTAL);
     bottomRow->AddStretchSpacer();
-    bottomRow->Add(back, 0, wxALL, 12);
+    bottomRow->Add(back, 0, wxRIGHT | wxLEFT | wxBOTTOM, 12);
 
     root->Add(bottomRow, 0, wxEXPAND);