mirror of
https://github.com/boykopovar/AnyPS5.git
synced 2026-10-18 23:07:07 -05:00
94 lines
4.3 KiB
C++
94 lines
4.3 KiB
C++
#include "PeImportFirst.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cctype>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <iterator>
|
|
#include <stdexcept>
|
|
|
|
namespace Allocator {
|
|
|
|
namespace {
|
|
|
|
std::uint32_t read32(const std::vector<unsigned char>& bytes, std::size_t offset) {
|
|
if (offset + 4 > bytes.size()) throw std::runtime_error("truncated image");
|
|
std::uint32_t value = 0;
|
|
std::memcpy(&value, bytes.data() + offset, sizeof(value));
|
|
return value;
|
|
}
|
|
|
|
std::uint16_t read16(const std::vector<unsigned char>& bytes, std::size_t offset) {
|
|
if (offset + 2 > bytes.size()) throw std::runtime_error("truncated image");
|
|
std::uint16_t value = 0;
|
|
std::memcpy(&value, bytes.data() + offset, sizeof(value));
|
|
return value;
|
|
}
|
|
|
|
std::string lower(std::string text) {
|
|
std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
|
return text;
|
|
}
|
|
|
|
}
|
|
|
|
void MoveImportFirst(std::vector<unsigned char>& bytes, const std::string& dll) {
|
|
const auto pe = read32(bytes, 0x3c);
|
|
if (read32(bytes, pe) != 0x00004550u) throw std::runtime_error("not a PE image");
|
|
const auto sections = read16(bytes, pe + 6);
|
|
const auto optionalSize = read16(bytes, pe + 20);
|
|
const auto optional = pe + 24;
|
|
if (read16(bytes, optional) != 0x20bu) throw std::runtime_error("not a PE32+ image");
|
|
const auto importRva = read32(bytes, optional + 112 + 8);
|
|
if (importRva == 0) throw std::runtime_error("the image has no imports");
|
|
const auto sectionTable = optional + optionalSize;
|
|
const auto fileOffset = [&](std::uint32_t rva) -> std::size_t {
|
|
for (std::uint16_t index = 0; index < sections; ++index) {
|
|
const auto header = sectionTable + index * 40u;
|
|
const auto virtualSize = read32(bytes, header + 8);
|
|
const auto virtualAddress = read32(bytes, header + 12);
|
|
const auto rawSize = read32(bytes, header + 16);
|
|
const auto rawOffset = read32(bytes, header + 20);
|
|
if (rva >= virtualAddress && rva < virtualAddress + std::max(virtualSize, rawSize)) {
|
|
if (rva - virtualAddress >= rawSize) throw std::runtime_error("RVA outside the section's file data");
|
|
return rawOffset + (rva - virtualAddress);
|
|
}
|
|
}
|
|
throw std::runtime_error("RVA outside every section");
|
|
};
|
|
const auto directory = fileOffset(importRva);
|
|
constexpr std::size_t DescriptorBytes = 20;
|
|
std::vector<std::array<unsigned char, DescriptorBytes>> descriptors;
|
|
std::size_t found = SIZE_MAX;
|
|
for (std::size_t offset = directory;; offset += DescriptorBytes) {
|
|
if (offset + DescriptorBytes > bytes.size()) throw std::runtime_error("truncated import directory");
|
|
std::array<unsigned char, DescriptorBytes> descriptor{};
|
|
std::memcpy(descriptor.data(), bytes.data() + offset, DescriptorBytes);
|
|
if (std::all_of(descriptor.begin(), descriptor.end(), [](unsigned char value) { return value == 0; })) break;
|
|
const auto nameOffset = fileOffset(read32(bytes, offset + 12));
|
|
std::string name;
|
|
for (auto at = nameOffset; at < bytes.size() && bytes[at] != 0; ++at) name.push_back(static_cast<char>(bytes[at]));
|
|
if (lower(name) == lower(dll)) found = descriptors.size();
|
|
descriptors.push_back(descriptor);
|
|
}
|
|
if (found == SIZE_MAX) throw std::runtime_error("the image does not import " + dll);
|
|
std::rotate(descriptors.begin(), descriptors.begin() + static_cast<std::ptrdiff_t>(found), descriptors.begin() + static_cast<std::ptrdiff_t>(found) + 1);
|
|
for (std::size_t index = 0; index < descriptors.size(); ++index) std::memcpy(bytes.data() + directory + index * DescriptorBytes, descriptors[index].data(), DescriptorBytes);
|
|
}
|
|
|
|
void MoveImportFirstInFile(const std::string& path, const std::string& dll) {
|
|
std::vector<unsigned char> bytes;
|
|
{
|
|
std::ifstream input(path, std::ios::binary);
|
|
if (!input) throw std::runtime_error("cannot open " + path);
|
|
bytes.assign(std::istreambuf_iterator<char>(input), std::istreambuf_iterator<char>());
|
|
}
|
|
MoveImportFirst(bytes, dll);
|
|
std::ofstream output(path, std::ios::binary | std::ios::trunc);
|
|
if (!output || !output.write(reinterpret_cast<const char*>(bytes.data()), static_cast<std::streamsize>(bytes.size()))) throw std::runtime_error("cannot write " + path);
|
|
}
|
|
|
|
}
|