/*!
 * \brief Find size of Linux kernel in boot image
 *
 * \pre The file position can be at any offset prior to calling this function.
 *
 * \post The file pointer position is undefined after this function returns.
 *       Use File::seek() to return to a known position.
 *
 * \param[in] reader Reader to set error message
 * \param[in] file File handle
 * \param[in] kernel_offset Offset of kernel in boot image
 * \param[out] kernel_size_out Pointer to store kernel size
 *
 * \return
 *   * Nothing if the kernel size is found
 *   * FileError::UnexpectedEof if the kernel size cannot be found
 *   * A specific error if any file operation fails
 */
oc::result<void>
LokiFormatReader::find_linux_kernel_size(Reader &reader,File &file,
                                         uint32_t kernel_offset,
                                         uint32_t &kernel_size_out)
{
    uint32_t kernel_size;

    // If the boot image was patched with an early version of loki, the
    // original kernel size is not stored in the loki header properly (or in the
    // shellcode). The size is stored in the kernel image's header though, so
    // we'll use that.
    // http://www.simtec.co.uk/products/SWLINUX/files/booting_article.html#d0e309
    auto seek_ret = file.seek(kernel_offset + 0x2c, SEEK_SET);
    if (!seek_ret) {
        if (file.is_fatal()) { reader.set_fatal(); }
        return seek_ret.as_failure();
    }

    auto ret = file_read_exact(file, &kernel_size, sizeof(kernel_size));
    if (!ret) {
        if (file.is_fatal()) { reader.set_fatal(); }
        return ret.as_failure();
    }

    kernel_size_out = mb_le32toh(kernel_size);
    return oc::success();
}
/*!
 * \brief Find and read Loki boot image header
 *
 * \note The integral fields in the header will be converted to the host's byte
 *       order.
 *
 * \pre The file position can be at any offset prior to calling this function.
 *
 * \post The file pointer position is undefined after this function returns.
 *       Use File::seek() to return to a known position.
 *
 * \param[in] reader Reader
 * \param[in] file File handle
 * \param[out] header_out Pointer to store header
 * \param[out] offset_out Pointer to store header offset
 *
 * \return
 *   * Nothing if the header is found
 *   * A LokiError if the header is not found
 *   * A specified error code if any file operation fails
 */
oc::result<void> LokiFormatReader::find_loki_header(Reader &reader, File &file,
                                                    LokiHeader &header_out,
                                                    uint64_t &offset_out)
{
    LokiHeader header;

    auto seek_ret = file.seek(LOKI_MAGIC_OFFSET, SEEK_SET);
    if (!seek_ret) {
        if (file.is_fatal()) { reader.set_fatal(); }
        return seek_ret.as_failure();
    }

    auto ret = file_read_exact(file, &header, sizeof(header));
    if (!ret) {
        if (ret.error() == FileError::UnexpectedEof) {
            return LokiError::LokiHeaderTooSmall;
        } else {
            if (file.is_fatal()) { reader.set_fatal(); }
            return ret.as_failure();
        }
    }

    if (memcmp(header.magic, LOKI_MAGIC, LOKI_MAGIC_SIZE) != 0) {
        return LokiError::InvalidLokiMagic;
    }

    loki_fix_header_byte_order(header);
    header_out = header;
    offset_out = LOKI_MAGIC_OFFSET;

    return oc::success();
}
static oc::result<void>
_mtk_compute_sha1(Writer &writer, SegmentWriter &seg,
                  File &file, unsigned char digest[SHA_DIGEST_LENGTH])
{
    SHA_CTX sha_ctx;
    char buf[10240];

    uint32_t kernel_mtkhdr_size = 0;
    uint32_t ramdisk_mtkhdr_size = 0;

    if (!SHA1_Init(&sha_ctx)) {
        return android::AndroidError::Sha1InitError;
    }

    for (auto const &entry : seg.entries()) {
        uint64_t remain = *entry.size;

        auto seek_ret = file.seek(static_cast<int64_t>(entry.offset),
                                  SEEK_SET);
        if (!seek_ret) {
            if (file.is_fatal()) { writer.set_fatal(); }
            return seek_ret.as_failure();
        }

        // Update checksum with data
        while (remain > 0) {
            auto to_read = std::min<uint64_t>(remain, sizeof(buf));

            auto ret = file_read_exact(file, buf, static_cast<size_t>(to_read));
            if (!ret) {
                if (writer.is_fatal()) { writer.set_fatal(); }
                return ret.as_failure();
            }

            if (!SHA1_Update(&sha_ctx, buf, static_cast<size_t>(to_read))) {
                return android::AndroidError::Sha1UpdateError;
            }

            remain -= to_read;
        }

        uint32_t le32_size;

        // Update checksum with size
        switch (entry.type) {
        case ENTRY_TYPE_MTK_KERNEL_HEADER:
            kernel_mtkhdr_size = *entry.size;
            continue;
        case ENTRY_TYPE_MTK_RAMDISK_HEADER:
            ramdisk_mtkhdr_size = *entry.size;
            continue;
        case ENTRY_TYPE_KERNEL:
            le32_size = mb_htole32(*entry.size + kernel_mtkhdr_size);
            break;
        case ENTRY_TYPE_RAMDISK:
            le32_size = mb_htole32(*entry.size + ramdisk_mtkhdr_size);
            break;
        case ENTRY_TYPE_SECONDBOOT:
            le32_size = mb_htole32(*entry.size);
            break;
        case ENTRY_TYPE_DEVICE_TREE:
            if (*entry.size == 0) {
                continue;
            }
            le32_size = mb_htole32(*entry.size);
            break;
        default:
            continue;
        }

        if (!SHA1_Update(&sha_ctx, &le32_size, sizeof(le32_size))) {
            return android::AndroidError::Sha1UpdateError;
        }
    }

    if (!SHA1_Final(digest, &sha_ctx)) {
        return android::AndroidError::Sha1UpdateError;
    }

    return oc::success();
}
/*!
 * \brief Find ramdisk size in old-style Loki image
 *
 * \pre The file position can be at any offset prior to calling this function.
 *
 * \post The file pointer position is undefined after this function returns.
 *       Use File::seek() to return to a known position.
 *
 * \param[in] reader Reader to set error message
 * \param[in] file File handle
 * \param[in] hdr Android header
 * \param[in] ramdisk_offset Offset of ramdisk in image
 * \param[out] ramdisk_size_out Pointer to store ramdisk size
 *
 * \return
 *   * Nothing if the ramdisk size is found
 *   * A LokiError if the ramdisk size is not found
 *   * A specific error if any file operation fails
 */
oc::result<void>
LokiFormatReader::find_ramdisk_size_old(Reader &reader, File &file,
                                        const android::AndroidHeader &hdr,
                                        uint32_t ramdisk_offset,
                                        uint32_t &ramdisk_size_out)
{
    int32_t aboot_size;

    // If the boot image was patched with an old version of loki, the ramdisk
    // size is not stored properly. We'll need to guess the size of the archive.

    // The ramdisk is supposed to be from the gzip header to EOF, but loki needs
    // to store a copy of aboot, so it is put in the last 0x200 bytes of the
    // file.
    if (is_lg_ramdisk_address(hdr.ramdisk_addr)) {
        aboot_size = static_cast<int32_t>(hdr.page_size);
    } else {
        aboot_size = 0x200;
    }

    auto aboot_offset = file.seek(-aboot_size, SEEK_END);
    if (!aboot_offset) {
        if (file.is_fatal()) { reader.set_fatal(); }
        return aboot_offset.as_failure();
    }

    if (ramdisk_offset > aboot_offset.value()) {
        return LokiError::RamdiskOffsetGreaterThanAbootOffset;
    }

    // Ignore zero padding as we might strip away too much
#if 1
    ramdisk_size_out = static_cast<uint32_t>(
            aboot_offset.value() - ramdisk_offset);
    return oc::success();
#else
    char buf[1024];

    // Search backwards to find non-zero byte
    uint64_t cur_offset = aboot_offset.value();

    while (cur_offset > ramdisk_offset) {
        size_t to_read = std::min<uint64_t>(
                sizeof(buf), cur_offset - ramdisk_offset);
        cur_offset -= to_read;

        auto seek_ret = file.seek(cur_offset, SEEK_SET);
        if (!seek_ret) {
            if (file.is_fatal()) { reader.set_fatal(); }
            return seek_ret.as_failure();
        }

        auto ret = file_read_exact(file, buf, to_read);
        if (!ret) {
            if (file.is_fatal()) { reader.set_fatal(); }
            return ret.as_failure();
        }

        for (size_t i = to_read; i-- > 0; ) {
            if (buf[i] != '\0') {
                ramdisk_size_out = cur_offset - ramdisk_offset + i;
                return oc::success();
            }
        }
    }

    return LokiError::FailedToDetermineRamdiskSize;
#endif
}
/*!
 * \brief Find gzip ramdisk offset in old-style Loki image
 *
 * This function will search for gzip headers (`0x1f8b08`) with a flags byte of
 * `0x00` or `0x08`. It will find the first occurrence of either magic string.
 * If both are found, the one with the flags byte set to `0x08` takes precedence
 * as it indiciates that the original filename field is set. This is usually the
 * case for ramdisks packed via the `gzip` command line tool.
 *
 * \pre The file position can be at any offset prior to calling this function.
 *
 * \post The file pointer position is undefined after this function returns.
 *       Use File::seek() to return to a known position.
 *
 * \param[in] reader Reader to set error message
 * \param[in] file File handle
 * \param[in] start_offset Starting offset for search
 * \param[out] gzip_offset_out Pointer to store gzip ramdisk offset
 *
 * \return
 *   * Nothing if a gzip offset is found
 *   * A LokiError if no gzip offsets are found
 *   * A specific error if any file operation fails
 */
oc::result<void>
LokiFormatReader::find_gzip_offset_old(Reader &reader, File &file,
                                       uint32_t start_offset,
                                       uint64_t &gzip_offset_out)
{
    struct SearchResult
    {
        std::optional<uint64_t> flag0_offset;
        std::optional<uint64_t> flag8_offset;
    };

    // gzip header:
    // byte 0-1 : magic bytes 0x1f, 0x8b
    // byte 2   : compression (0x08 = deflate)
    // byte 3   : flags
    // byte 4-7 : modification timestamp
    // byte 8   : compression flags
    // byte 9   : operating system

    static const unsigned char gzip_deflate_magic[] = { 0x1f, 0x8b, 0x08 };

    SearchResult result = {};

    // Find first result with flags == 0x00 and flags == 0x08
    auto result_cb = [&](File &file_, uint64_t offset)
            -> oc::result<FileSearchAction> {
        unsigned char flags;

        // Stop early if possible
        if (result.flag0_offset && result.flag8_offset) {
            return FileSearchAction::Stop;
        }

        // Save original position
        OUTCOME_TRY(orig_offset, file_.seek(0, SEEK_CUR));

        // Seek to flags byte
        OUTCOME_TRYV(file_.seek(static_cast<int64_t>(offset + 3), SEEK_SET));

        // Read next bytes for flags
        auto ret = file_read_exact(file_, &flags, sizeof(flags));
        if (!ret) {
            if (ret.error() == FileError::UnexpectedEof) {
                return FileSearchAction::Stop;
            } else {
                return ret.as_failure();
            }
        }

        if (!result.flag0_offset && flags == 0x00) {
            result.flag0_offset = offset;
        } else if (!result.flag8_offset && flags == 0x08) {
            result.flag8_offset = offset;
        }

        // Restore original position as per contract
        OUTCOME_TRYV(file_.seek(static_cast<int64_t>(orig_offset), SEEK_SET));

        return FileSearchAction::Continue;
    };

    auto ret = file_search(file, start_offset, {}, 0, gzip_deflate_magic,
                           sizeof(gzip_deflate_magic), {}, result_cb);
    if (!ret) {
        if (file.is_fatal()) { reader.set_fatal(); }
        return ret.as_failure();
    }

    // Prefer gzip header with original filename flag since most loki'd boot
    // images will have been compressed manually with the gzip tool
    if (result.flag8_offset) {
        gzip_offset_out = *result.flag8_offset;
    } else if (result.flag0_offset) {
        gzip_offset_out = *result.flag0_offset;
    } else {
        return LokiError::NoRamdiskGzipHeaderFound;
    }

    return oc::success();
}
/*!
 * \brief Find and read Loki ramdisk address
 *
 * \pre The file position can be at any offset prior to calling this function.
 *
 * \post The file pointer position is undefined after this function returns.
 *       Use File::seek() to return to a known position.
 *
 * \param[in] reader Reader to set error message
 * \param[in] file File handle
 * \param[in] hdr Android header
 * \param[in] loki_hdr Loki header
 * \param[out] ramdisk_addr_out Pointer to store ramdisk address
 *
 * \return
 *   * Nothing if the ramdisk address is found
 *   * A LokiError if the ramdisk address is not found
 *   * A specific error code if any file operation fails
 */
oc::result<void>
LokiFormatReader::find_ramdisk_address(Reader &reader, File &file,
                                       const android::AndroidHeader &hdr,
                                       const LokiHeader &loki_hdr,
                                       uint32_t &ramdisk_addr_out)
{
    // If the boot image was patched with a newer version of loki, find the
    // ramdisk offset in the shell code
    uint32_t ramdisk_addr = 0;

    if (loki_hdr.ramdisk_addr != 0) {
        uint64_t offset = 0;

        auto result_cb = [&](File &file_, uint64_t offset_)
                -> oc::result<FileSearchAction> {
            (void) file_;
            offset = offset_;
            return FileSearchAction::Continue;
        };

        auto ret = file_search(file, {}, {}, 0, LOKI_SHELLCODE,
                               LOKI_SHELLCODE_SIZE - 9, 1, result_cb);
        if (!ret) {
            if (file.is_fatal()) { reader.set_fatal(); }
            return ret.as_failure();
        }

        if (offset == 0) {
            return LokiError::ShellcodeNotFound;
        }

        offset += LOKI_SHELLCODE_SIZE - 5;

        auto seek_ret = file.seek(static_cast<int64_t>(offset), SEEK_SET);
        if (!seek_ret) {
            if (file.is_fatal()) { reader.set_fatal(); }
            return seek_ret.as_failure();
        }

        auto read_ret = file_read_exact(file, &ramdisk_addr,
                                        sizeof(ramdisk_addr));
        if (!read_ret) {
            if (file.is_fatal()) { reader.set_fatal(); }
            return read_ret.as_failure();
        }

        ramdisk_addr = mb_le32toh(ramdisk_addr);
    } else {
        // Otherwise, use the default for jflte (- 0x00008000 + 0x02000000)

        if (hdr.kernel_addr > UINT32_MAX - 0x01ff8000) {
            //DEBUG("Invalid kernel address: %" PRIu32, hdr.kernel_addr);
            return LokiError::InvalidKernelAddress;
        }

        ramdisk_addr = hdr.kernel_addr + 0x01ff8000;
    }

    ramdisk_addr_out = ramdisk_addr;
    return oc::success();
}