Skip to content

Instantly share code, notes, and snippets.

@Redchards
Last active March 23, 2016 11:07
Show Gist options
  • Select an option

  • Save Redchards/d8fb32691e5db852ac7f to your computer and use it in GitHub Desktop.

Select an option

Save Redchards/d8fb32691e5db852ac7f to your computer and use it in GitHub Desktop.
Sample taken from my PeParser project to be reused.
#pragma once
#error This file need modification on linux, be aware of this. Once you made it, you can delete this error.
#include <cstdint>
#include <Windows.h>
#include <winnt.h>
#ifdef DLL_EXPORT
# define PE_API __declspec(dllexport)
#else
# define PE_API __declspec(dllimport)
#endif
#ifdef __WIN64
# define X64
#else
# define x86
#endif
#ifdef X86
typedef unsigned long size_type;
#else
typedef size_t size_type;
#endif
typedef uint16_t flag_type;
typedef DWORD image_data_type;
#define stringify(x) stringify_(x)
#define stringify_(x) #x
enum class Endianess : flag_type
{
little,
big
};
#pragma once
#include <Configuration.h>
#include <fstream>
#include <iterator>
#include <string>
#include <unordered_map>
#include <vector>
enum class StreamGoal : flag_type
{
read,
write
};
// Selector, allowing us to add new class easily, along with stream goals
template<StreamGoal goal>
struct FileStreamSelector;
template<>
class FileStreamSelector<StreamGoal::read>
{
public:
using type = std::ifstream;
public:
FileStreamSelector(const std::string& filename, std::ios::ios_base::openmode flags)
{
fstream_ = &streamMap_[filename];
if (!fstream_->is_open())
{
fstream_->open(filename, flags);
}
}
void read(char* buffer, std::streampos position, size_type size)
{
goTo(position);
fstream_->read(buffer, size);
}
void read(unsigned char* buffer, std::streampos position, size_type size)
{
read(reinterpret_cast<char*>(buffer), position, size);
}
template<size_type n>
void read(std::array<char, n>& buffer, std::streampos position, size_type size)
{
read(buffer.data(), position, size);
}
template<size_type n>
void read(std::array<unsigned char, n>& buffer, std::streampos position, size_type size)
{
read(buffer.data(), position, size);
}
void read(std::vector<char>& buffer, std::streampos position, size_type size)
{
read(buffer.data(), position, size);
}
void read(std::vector<unsigned char>& buffer, std::streampos position, size_type size)
{
read(buffer.data(), position, size);
}
void rewind()
{
goTo(0);
}
void goTo(std::streampos position)
{
if (!fstream_->seekg(position))
{
throw std::ios_base::failure("Error when processing the file !");
}
}
std::streampos getCurrentPosition()
{
return fstream_->tellg();
}
protected:
type* fstream_;
static std::unordered_map<std::string, type> streamMap_;
};
template<>
class FileStreamSelector<StreamGoal::write>
{
public:
using type = std::ofstream;
public:
FileStreamSelector(const std::string& filename, std::ios::ios_base::openmode flags)
{
fstream_ = &streamMap_[filename];
if (!fstream_->is_open())
{
fstream_->open(filename, flags);
}
}
void write(const char* buffer, std::streampos position, size_type size)
{
goTo(position);
fstream_->write(buffer, size);
}
void write(const unsigned char* buffer, std::streampos position, size_type size)
{
write(reinterpret_cast<const char*>(buffer), position, size);
}
template<size_type n>
void write(const std::array<char, n>& buffer, std::streampos position, size_type size)
{
write(buffer.data(), position, size);
}
template<size_type n>
void write(const std::array<unsigned char, n>& buffer, std::streampos position, size_type size)
{
write(buffer.data(), position, size);
}
void write(const std::vector<char>& buffer, std::streampos position, size_type size)
{
write(buffer.data(), position, size);
}
void write(const std::vector<unsigned char>& buffer, std::streampos position, size_type size)
{
write(buffer.data(), position, size);
}
void rewind()
{
goTo(0);
}
void goTo(std::streampos position)
{
if (!fstream_->seekp(position))
{
throw std::ios_base::failure("Error when processing the file !");
}
}
std::streampos getCurrentPosition()
{
return fstream_->tellp();
}
protected:
type* fstream_;
static std::unordered_map<std::string, type> streamMap_;
};
template<StreamGoal goal>
class FileStreamBase : public FileStreamSelector<goal>
{
protected:
FileStreamBase(const std::string& filename, std::ios::ios_base::openmode flags) : FileStreamSelector(filename, flags),
filename_(filename)
{
init();
}
public:
void loadFile(const std::string& filename, std::ios::ios_base::openmode flags = std::ios_base::in)
{
unloadFile();
fstream_->open(filename, flags);
filename_ = filename;
init();
}
void unloadFile()
{
if (fstream_->is_open())
{
fstream_->close();
filename_ = "";
}
}
bool isOpen()
{
return fstream_->is_open();
}
const std::string& getCurrentFileName() const noexcept
{
return filename_;
}
private:
void init()
{
if (!fstream_->good())
{
throw std::ios_base::failure((std::string("Error : failed to open the file ") + getCurrentFileName() + ". Please check that the file exists, and is a valid Win32 file !").c_str());
}
// So that stream will not interpret "whitespace" bytes as real whitespaces, and remove them.
fstream_->unsetf(std::ios::skipws);
}
private:
std::string filename_;
};
#include <stdafx.h>
#include <FileValueReader.h>
// OPTIMIZATION NOTE :
// We can have one reader by file (not multiplicating readers), by having a static undordered_map of readers, and doing a lookup each time we want to access
// a file. Supposedly, the lookup will be negated by the read, but this requires benchmark.
FileValueReaderBase::FileValueReaderBase(const std::string& filename, std::ios_base::openmode flags) : Base(filename, flags)
{}
// TODO : Remove maybe ?
std::unique_ptr<unsigned char> FileValueReaderBase::retrieveRawData(std::streampos position, size_type size)
{
std::unique_ptr<unsigned char> dataPtr(new unsigned char[size]);
read(dataPtr.get(), position, size);
return dataPtr;
}
std::vector<unsigned char> FileValueReaderBase::retrieveRawBuffer(std::streampos position, size_type size)
{
std::vector<unsigned char> buffer;
buffer.resize(size);
read(buffer, position, size);
return buffer;
}
#pragma once
#include <Configuration.h>
#include <FileStream.h>
#include <array>
#include <fstream>
#include <iterator>
#include <memory>
#include <string>
#include <vector>
// Maybe should move all enums to Configuration.h header
// NOTE : This is the endianess of the file as one might want to analyze a file with a different endianess.
// However, PE/COFF files should always be little-endian (Windows is only present on little-endian platforms, except WinCE)
enum class BufferPolicy : flag_type
{
nonBuffered = false,
buffered = true
};
class PE_API FileValueReaderBase : public FileStreamBase<StreamGoal::read>
{
using Base = FileStreamBase<StreamGoal::read>;
protected:
FileValueReaderBase(const std::string& filename, std::ios::ios_base::openmode flags = std::ios_base::in | std::ios::binary);
public:
std::unique_ptr<uint8_t> retrieveRawData(std::streampos position, size_type size);
std::vector<uint8_t> retrieveRawBuffer(std::streampos position, size_type size);
};
template<Endianess endian = Endianess::little>
class FileValueReader;
template<>
class PE_API FileValueReader<Endianess::little> : public FileValueReaderBase
{
public:
FileValueReader(const std::string& filename, std::ios::ios_base::openmode flags = std::ios_base::in | std::ios::binary) : FileValueReaderBase(filename, flags)
{}
size_type retrieveValue(std::streampos position, size_type size)
{
if (size > sizeof(size_type))
{
// Throwing necessary here ?
throw std::ios_base::failure("Error : attempt to read to much bytes at once. Check the maximal value of a field on your platform");
}
std::array<unsigned char, sizeof(size_type)> value{};
size_type out = 0;
read(value, position, size);
for (size_type i = size; i != 0; --i)
{
out |= (value[i - 1] << ((i - 1) * 8));
}
return out;
}
};
template<>
class PE_API FileValueReader<Endianess::big> : public FileValueReaderBase
{
public:
FileValueReader(const std::string& filename, std::ios::ios_base::openmode flags = std::ios_base::in | std::ios::binary) : FileValueReaderBase(filename, flags)
{}
size_type retrieveValue(std::streampos position, size_type size)
{
if (size > sizeof(size_type))
{
throw std::ios_base::failure("Error : attempt to read to much bytes at once. Check the maximal value of a field on your platform");
}
std::array<unsigned char, sizeof(size_type)> value{};
size_type out = 0;
read(value, position, size);
for (size_type i = 0; i < size; ++i)
{
out |= (value[i] << (i * 8));
}
return out;
}
};
using FileReader = FileValueReader<Endianess::little>;
#include <stdafx.h>
#include <FileValueWriter.h>
FileValueWriterBase::FileValueWriterBase(const std::string& filename, std::ios::ios_base::openmode flags) : Base(filename, flags)
{}
void FileValueWriterBase::checkWrite() const
{
if (!fstream_->good())
{
throw std::ios_base::failure(std::string("Error when writing data to the file '") + getCurrentFileName() + "' !");
}
}
void FileValueWriterBase::writeRawDataAt(const std::vector<char>& data, std::streampos position)
{
write(data, position, data.size());
checkWrite();
}
void FileValueWriterBase::writeRawDataAt(const char* const data, std::streamsize size, std::streampos position)
{
write(data, position, size);
checkWrite();
}
#pragma once
#include <Configuration.h>
#include <FileStream.h>
#include <fstream>
#include <string>
#include <vector>
class FileValueWriterBase : public FileStreamBase<StreamGoal::write>
{
using Base = FileStreamBase<StreamGoal::write>;
protected:
FileValueWriterBase(const std::string& filename, std::ios::ios_base::openmode flags = std::ios_base::out);
void checkWrite() const;
public:
void writeRawDataAt(const std::vector<char>& data, std::streampos position);
void writeRawDataAt(const char* const data, std::streamsize size, std::streampos position);
};
template<Endianess endian = Endianess::little>
class FileValueWriter;
template<>
class FileValueWriter<Endianess::little> : public FileValueWriterBase
{
public:
FileValueWriter(const std::string& filename, std::ios::ios_base::openmode flags = std::ios_base::out) : FileValueWriterBase(filename, flags)
{}
void writerValue(const std::vector<char>& data, std::streampos position)
{
}
void writeValue(const char* const data, size_type size, std::streampos position)
{
}
};
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment