/*
 * palera1n - https://palera.in
 *
 * Copyright (C) 2026 palera1n team
 *
 * SPDX-License-Identifier: MIT
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 */

#include "exploit.h"

#include <stdio.h>
#include <stdlib.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <pthread.h>
#include <stdint.h>
#include <unistd.h>
#include <string.h>

#include "globals.h"
#include "paleinfo.h"
#include "pongo_helper.h"

#if WITH_CIDERRAIN
# include <ciderra1n/log.h>
# include <ciderra1n/common.h>
# include <ciderra1n/usb.h>
# include <ciderra1n/ra1n.h>
# include <ciderra1n/dfu.h>
#else
# include <openra1n/utils.h>
# include <openra1n/checkm8.h>
# include <openra1n/dfu.h>
# include <openra1n/payload.h>
# include <openra1n/shim.h>
#endif

#define VID_APPLE   (0x5AC)
#define PID_PONGO   (0x4141)
#define PID_DFU     (0x1227)

const char* stage_to_string(int stage)
{
    switch (stage) {
        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!";
        case STAGE_ERROR:       return "Error!";
        default:                return "Waiting...";
    }
}

int stage_to_progress(int stage)
{
    switch (stage) {
        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;
        case STAGE_ERROR:       return 99;
        default:                return 0;
    }
}

const char* err_to_string(int err) {
    // error codes should be the same between both checkm8 backends
    switch (err) {
        case 0:     return "";
        case 1:     return "-p1_client_err_abort_success";
        case 2:     return "-p1_client_err_unsupported";
        case 3:     return "-p1_client_err_response_size";
        case 4:     return "-p1_client_err_abort";
        case 5:     return "-p1_client_err_timeout";
        case 6:     return "-p1_client_err_pipe_stalled";
        case 7:     return "-p1_client_err_no_device";
        case 8:     return "-p1_client_err_notresponding";
        case 9:     return "-p1_client_err_invalid_response";
        case 10:    return "-p1_client_err_async_download_failed";
        case 11:    return "-p1_client_err_transaction_abort_failed";
        case 12:    return "-p1_client_err_status";
        case 13:    return "-p1_client_err_not_dfu_idle";
        case 14:    return "-p1_client_err_no_stage1_generator";
        case 15:    return "-p1_client_err_invalid_payload";
        case 16:    return "-p1_client_err_no_image";
        case 17:    return "-p1_client_err_not_dfu";
        case 18:    return "-p1_client_err_soft_dfu";
        case 19:    return "-p1_client_err_cnt_overflow";
        default:    return "-p1_client_err_undefined";
    }
}

#if WITH_CIDERRAIN

// MARK: exploit (libciderra1n)

// NOTE: this backend is NOT publicly available
//
// ciderra1n is based on the checkra1n 1337 implementation
// of the checkm8 exploit using DFU aborts instead of usb resets
//
// this supports Linux (liteusb), and Darwin (IOKit)
//
// this also fixes a lot of issues with AMD devices.
//
// this also supports A7 devices.
//
// windows eta s0n?

int exploit(int *stage)
{
    LOG("Waiting for devices...");

    *stage = STAGE_RESET;

    client_t* client = NULL;
    if (!alloc_usb_client(&client)) {
        *stage = STAGE_ERROR;
        return RA1N_CLIENT_ERR_UNDEFINED;
    }

    ra1n_client_err_t cr = RA1N_CLIENT_ERR_NONE;

    uint32_t flag = 0;
    uint16_t cpid = 0;

    for (;;) {
        if (usb_open(client) != 0) {
            // no usb device
            // slight delay is required for AMD
            usleep(500000);
            goto reloop;
        }

        if (!usb_get_client_flag(client, &flag))
            goto reloop;

        if (flag & kDeviceDFUModeFlag) {
            if (flag & kDeviceStandardPwnedDFUModeFlag) {
                // we cant do shit sorry
                LOG_ERROR("Found an unknown pwned DFU mode device");
                usb_close(client);
                *stage = STAGE_ERROR;
                continue; // err
            }

            if (flag & kDevicePwnedDFUModeFlag) {
                // we cant do shit sorry
                LOG_ERROR("Detected pwned DFU mode device, please enter normal DFU mode.");
                usb_close(client);
                *stage = STAGE_ERROR;
                continue; // err
            }

            if (!(flag & kDeviceNotSoftDFUModeFlag)) {
                // we cant do shit sorry
                LOG_ERROR("Found soft-DFU mode device");
                usb_close(client);
                *stage = STAGE_ERROR;
                break; // err
            }

            LOG("Found DFU mode device");

            if (*stage == STAGE_RESET) {
                cr = ra1n_stage_reset(client);
                if (cr == RA1N_CLIENT_ERR_NONE || cr == RA1N_CLIENT_ERR_ABORT_SUCCESS || cr == RA1N_CLIENT_ERR_NOTRESPONDING) {
                    usb_reset(client, kUSBResetDevice | kUSBReEnumerate);
                    usb_close(client);
                    if (cr == RA1N_CLIENT_ERR_NONE) { // success
                        *stage = STAGE_SETUP;
                    }
                    continue;
                }
                LOG_ERROR("STAGE_RESET failed (rv=%d, reason=%s)", cr, err_to_string(cr));
                usb_close(client);
                *stage = STAGE_ERROR;
                break; // err
            } else if (*stage == STAGE_SETUP) {
                cr = ra1n_stage_setup(client);
                if (cr == RA1N_CLIENT_ERR_NONE) { // success
                    usb_reset(client, kUSBResetDevice | kUSBReEnumerate);
                    usb_close(client);
                    *stage = STAGE_PATCH;
                    continue;
                }
                LOG_ERROR("STAGE_SETUP failed (rv=%d, reason=%s)", cr, err_to_string(cr));
                usb_close(client);
                *stage = STAGE_ERROR;
                break; // err
            } else if (*stage == STAGE_PATCH) {
                cr = ra1n_stage_trigger(client, NULL);
                if (cr == RA1N_CLIENT_ERR_NONE) { // exploited
                    usb_reset(client, kUSBResetDevice | kUSBReEnumerate);
                    usb_close(client);
                    *stage = STAGE_YOLODFU;
                    LOG("Checkmate?");
                    continue;
                }
                LOG_ERROR("STAGE_PATCH failed (rv=%d, reason=%s)", cr, err_to_string(cr));
                usb_close(client);
                *stage = STAGE_ERROR;
                break; // err
            } else if (*stage == STAGE_YOLODFU) {
                LOG_ERROR("STAGE_YOLODFU failed (rv=%d, reason=%s)", cr, err_to_string(cr));
                usb_close(client);
                *stage = STAGE_ERROR;
                break; // err
            }
            continue;
        } else if (flag & kDeviceYoloDFUModeFlag) {
            LOG("Found download mode device");

            *stage = STAGE_YOLODFU;
            cr = send_compressed_pongo(client, g_payload_pongo.data, g_payload_pongo.data_len); // yoloDFU -> pongoOS
            if (cr == RA1N_CLIENT_ERR_NONE) {
                usb_reset(client, kUSBResetDevice | kUSBReEnumerate);
                usb_close(client); // success
                continue;
            }
            LOG_ERROR("Failed to send pongo image (rv=%d, reason=%s)", cr, err_to_string(cr));
            usb_close(client);
            *stage = STAGE_ERROR;
            break; // err
        } else if (flag & kDevicePongoModeFlag) {
            LOG("Found pongo mode device");

            if (palerain_flags & palerain_option_pongo_exit) {
                usb_close(client);
                *stage = STAGE_DONE;
                break; // success
            } else {
                LOG("Booting...");
                cr = send_full_pongo_jailbreak(client);
                if (cr != RA1N_CLIENT_ERR_NONE) {
                    LOG_ERROR("Failed to boot device (rv=%d, reason=%s)", cr, err_to_string(cr));
                    usb_close(client);
                    *stage = STAGE_ERROR;
                    break; // err
                }

                // assume everything went well, though
                // depending on the components sent to
                // the device, the device may not even
                // boot at all (not our fault xD)
                *stage = STAGE_DONE;
                break; // success
            }

            usb_close(client);
            continue;
        } else {
            LOG_ERROR("Found unknown device");
            usb_close(client);
            *stage = STAGE_ERROR;
            break; // err
        }

    reloop:
        usb_close(client);
    }

    return cr;
}

#else // WITH_CIDERRAIN

// MARK: exploit (libopenra1n)

// NOTE: this backend is publicly available
//
// openra1n is based on the gaster implementation of the exploit
// which is the more well-known implementation of checkm8
//
// supports Windows (Libusb), Linux (Libusb), and Darwin (IOKit)
//
// does not support AMD devices, AMD mobile will probably work though.
//
// on the contrary while gaster *does* support A7, openra1n does
// not. (yet to be implemented, cc @Mineek)
//
// Windows is technically *supported*, but during or before the
// exploit process libusbK drivers will need to be installed due
// to needing USB resets. However, this breaks other things within
// the TUI, GUI (... I think anyway) and itunes itself, and requires
// the drivers to be uninstalled afterwards for standard usage again.

// TODO: add support for A7 devices, and AMD devices
// TODO: libusbK driver installation?

int exploit(int *stage)
{
    LOG("Waiting for devices...");

    struct DeviceConfiguration deviceConfig;
    struct PayloadConfiguration payloadConfig;

    *stage = STAGE_RESET;
    checkm8_err_t cr = CHECKM8_ERR_NONE;

    usb_handle_t dfuHandle;
    usb_handle_t pongoHandle;

    init_usb_handle(&dfuHandle, VID_APPLE, PID_DFU);
    init_usb_handle(&pongoHandle, VID_APPLE, PID_PONGO);

    for (;;) {
        // pongoOS is checked first
        if (wait_usb_handle(&pongoHandle)) {
            LOG("Found pongo mode device");
            *stage = STAGE_PONGO;

            if (palerain_flags & palerain_option_pongo_exit)
                goto pongo_done;

            cr = send_full_pongo_jailbreak(&pongoHandle);
            if (cr != CHECKM8_ERR_NONE) {
                LOG_ERROR("Failed to boot device (rv=%d, reason=%s)", cr, err_to_string(cr));
                *stage = STAGE_ERROR;
                goto cleanup;
            }

            // assume everything went well, though
            // depending on the components sent to
            // the device, the device may not even
            // boot at all (not our fault xD)
        pongo_done:
            *stage = STAGE_DONE;
            goto cleanup;
        }

        // DFU is checked second
        if (wait_usb_handle(&dfuHandle)) {
            char *sn = get_usb_serial_number(&dfuHandle);
            if (!sn) goto reloop;

            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);
                goto reloop;
            }

            if (dfu_serial_number_is_pwned(sn)) {
                // we cant do shit sorry
                LOG_ERROR("Detected pwned DFU mode device, please enter normal DFU mode.");
                free(sn);
                *stage = STAGE_ERROR;
                goto cleanup;
            }

            if (dfu_serial_number_is_in_dfu_mode(sn)) {
                LOG("Found DFU mode device");
            } else if (dfu_serial_number_is_in_yolo_dfu(sn)) {
                LOG("Found download mode device");
                *stage = STAGE_YOLODFU;
            } else {
                free(sn);
                goto reloop;
            }

            free(sn);

            switch (*stage) {
                case STAGE_RESET:
                    cr = checkm8_stage_reset(&dfuHandle);
                    if (cr != CHECKM8_ERR_NONE) { // err
                        LOG_ERROR("STAGE_RESET failed (rv=%d, reason=%s)", cr, err_to_string(cr));
                        *stage = STAGE_RESET;
                        goto reloop;
                    }

                    *stage = STAGE_SETUP;
                    goto reloop; // success
                case STAGE_SETUP:
                    cr = checkm8_stage_setup(&dfuHandle, &deviceConfig);
                    if (cr != CHECKM8_ERR_NONE) { // err
                        LOG_ERROR("STAGE_SETUP failed (rv=%d, reason=%s)", cr, err_to_string(cr));
                        *stage = STAGE_RESET;
                        goto reloop;
                    }

                    *stage = STAGE_SPRAY;
                    goto reloop;
                case STAGE_SPRAY:
                    cr = checkm8_stage_spray(&dfuHandle, &deviceConfig);
                    if (cr != CHECKM8_ERR_NONE) { // err
                        LOG_ERROR("STAGE_SPRAY failed (rv=%d, reason=%s)", cr, err_to_string(cr));
                        *stage = STAGE_RESET;
                        goto reloop;
                    }

                    *stage = STAGE_PATCH;
                    goto reloop; // success
                case STAGE_PATCH:
                    cr = checkm8_stage_patch(&dfuHandle, &deviceConfig, &payloadConfig);
                    if (cr != CHECKM8_ERR_NONE) // err
                        LOG_ERROR("STAGE_PATCH failed (rv=%d, reason=%s)", cr, err_to_string(cr));

                    *stage = STAGE_RESET;
                    LOG("Checkmate?");
                    goto reloop; // success
                case STAGE_YOLODFU:
                    cr = send_compressed_pongo(&dfuHandle, g_payload_pongo.data, g_payload_pongo.data_len);
                    if (cr != CHECKM8_ERR_NONE) {
                        LOG_ERROR("Failed to send pongo image (rv=%d, reason=%s)", cr, err_to_string(cr));
                        *stage = STAGE_ERROR;
                        goto cleanup; // err
                    };

                    continue; // success
            }

        reloop:
            reset_usb_handle(&dfuHandle);
            continue;
        }
    }

cleanup:
    close_usb_handle(&dfuHandle);
    close_usb_handle(&pongoHandle);
    return cr;
}

#endif // WITH_CIDERRAIN