Add API for registering transforms

2 files changed, 357 insertions, 0 deletions

diff --git a/binaryninjaapi.cpp b/binaryninjaapi.cpp
index 026e93b3..677c7c27 100644
--- a/binaryninjaapi.cpp
+++ b/binaryninjaapi.cpp
@@ -773,6 +773,21 @@ void BinaryReader::Read(void* dest, size_t len)
 }
 
 
+DataBuffer BinaryReader::Read(size_t len)
+{
+	DataBuffer result(len);
+	Read(result.GetData(), len);
+	return result;
+}
+
+
+string BinaryReader::ReadString(size_t len)
+{
+	DataBuffer result = Read(len);
+	return string((const char*)result.GetData(), len);
+}
+
+
 uint8_t BinaryReader::Read8()
 {
 	uint8_t result;
@@ -869,6 +884,23 @@ bool BinaryReader::TryRead(void* dest, size_t len)
 }
 
 
+bool BinaryReader::TryRead(DataBuffer& dest, size_t len)
+{
+	dest.SetSize(len);
+	return TryRead(dest.GetData(), len);
+}
+
+
+bool BinaryReader::TryReadString(string& dest, size_t len)
+{
+	DataBuffer result(len);
+	if (!TryRead(result.GetData(), len))
+		return false;
+	dest = string((const char*)result.GetData(), len);
+	return true;
+}
+
+
 bool BinaryReader::TryRead8(uint8_t& result)
 {
 	return BNRead8(m_stream, &result);
@@ -985,6 +1017,18 @@ void BinaryWriter::Write(const void* src, size_t len)
 }
 
 
+void BinaryWriter::Write(const DataBuffer& buf)
+{
+	Write(buf.GetData(), buf.GetLength());
+}
+
+
+void BinaryWriter::Write(const string& str)
+{
+	Write(str.c_str(), str.size());
+}
+
+
 void BinaryWriter::Write8(uint8_t val)
 {
 	if (!BNWrite8(m_stream, val))
@@ -1061,6 +1105,18 @@ bool BinaryWriter::TryWrite(const void* src, size_t len)
 }
 
 
+bool BinaryWriter::TryWrite(const DataBuffer& buf)
+{
+	return TryWrite(buf.GetData(), buf.GetLength());
+}
+
+
+bool BinaryWriter::TryWrite(const string& str)
+{
+	return TryWrite(str.c_str(), str.size());
+}
+
+
 bool BinaryWriter::TryWrite8(uint8_t val)
 {
 	return BNWrite8(m_stream, val);
@@ -1137,3 +1193,241 @@ void BinaryWriter::SeekRelative(int64_t offset)
 {
 	BNSeekBinaryWriterRelative(m_stream, offset);
 }
+
+
+Transform::Transform(BNTransform* xform): m_xform(xform)
+{
+}
+
+
+Transform::Transform(BNTransformType type, const string& name, const string& longName):
+	m_xform(nullptr), m_typeForRegister(type), m_nameForRegister(name), m_longNameForRegister(longName)
+{
+}
+
+
+BNTransformParameterInfo* Transform::GetParametersCallback(void* ctxt, size_t* count)
+{
+	Transform* xform = (Transform*)ctxt;
+	vector<TransformParameter> params = xform->GetParameters();
+	*count = params.size();
+	BNTransformParameterInfo* result = new BNTransformParameterInfo[params.size()];
+
+	for (size_t i = 0; i < params.size(); i++)
+	{
+		result[i].name = BNAllocString(params[i].name.c_str());
+		result[i].longName = BNAllocString(params[i].longName.c_str());
+		result[i].fixedLength = params[i].fixedLength;
+	}
+
+	return result;
+}
+
+
+void Transform::FreeParametersCallback(BNTransformParameterInfo* params, size_t count)
+{
+	for (size_t i = 0; i < count; i++)
+	{
+		BNFreeString(params[i].name);
+		BNFreeString(params[i].longName);
+	}
+	delete[] params;
+}
+
+
+bool Transform::DecodeCallback(void* ctxt, BNDataBuffer* input, BNDataBuffer* output, BNTransformParameter* params, size_t paramCount)
+{
+	map<string, DataBuffer> paramMap;
+	for (size_t i = 0; i < paramCount; i++)
+		paramMap[params[i].name] = DataBuffer(BNDuplicateDataBuffer(params[i].value));
+
+	DataBuffer inputBuf(BNDuplicateDataBuffer(input));
+	DataBuffer outputBuf;
+
+	Transform* xform = (Transform*)ctxt;
+	bool result = xform->Decode(inputBuf, outputBuf, paramMap);
+	BNAssignDataBuffer(output, outputBuf.GetBufferObject());
+	return result;
+}
+
+
+bool Transform::EncodeCallback(void* ctxt, BNDataBuffer* input, BNDataBuffer* output, BNTransformParameter* params, size_t paramCount)
+{
+	map<string, DataBuffer> paramMap;
+	for (size_t i = 0; i < paramCount; i++)
+		paramMap[params[i].name] = DataBuffer(BNDuplicateDataBuffer(params[i].value));
+
+	DataBuffer inputBuf(BNDuplicateDataBuffer(input));
+	DataBuffer outputBuf;
+
+	Transform* xform = (Transform*)ctxt;
+	bool result = xform->Encode(inputBuf, outputBuf, paramMap);
+	BNAssignDataBuffer(output, outputBuf.GetBufferObject());
+	return result;
+}
+
+
+vector<TransformParameter> Transform::EncryptionKeyParameters(size_t fixedKeyLength)
+{
+	vector<TransformParameter> params;
+	TransformParameter key;
+	key.name = "key";
+	key.longName = "Encryption key";
+	key.fixedLength = fixedKeyLength;
+	params.push_back(key);
+	return params;
+}
+
+
+vector<TransformParameter> Transform::EncryptionKeyAndIVParameters(size_t fixedKeyLength, size_t fixedIVLength)
+{
+	vector<TransformParameter> params;
+	TransformParameter key, iv;
+	key.name = "key";
+	key.longName = "Encryption key";
+	key.fixedLength = fixedKeyLength;
+	iv.name = "iv";
+	iv.longName = "Initialization vector";
+	iv.fixedLength = fixedIVLength;
+	params.push_back(key);
+	params.push_back(iv);
+	return params;
+}
+
+
+void Transform::Register(Transform* xform)
+{
+	BNCustomTransform callbacks;
+	callbacks.context = xform;
+	callbacks.getParameters = GetParametersCallback;
+	callbacks.freeParameters = FreeParametersCallback;
+	callbacks.decode = DecodeCallback;
+	callbacks.encode = EncodeCallback;
+	xform->m_xform = BNRegisterTransformType(xform->m_typeForRegister, xform->m_nameForRegister.c_str(),
+		xform->m_longNameForRegister.c_str(), &callbacks);
+}
+
+
+Ref<Transform> Transform::GetByName(const string& name)
+{
+	BNTransform* result = BNGetTransformByName(name.c_str());
+	if (!result)
+		return nullptr;
+	return new CoreTransform(result);
+}
+
+
+vector<Ref<Transform>> Transform::GetTransformTypes()
+{
+	size_t count;
+	BNTransform** list = BNGetTransformTypeList(&count);
+
+	vector<Ref<Transform>> result;
+	for (size_t i = 0; i < count; i++)
+		result.push_back(new CoreTransform(BNNewTransformReference(list[i])));
+
+	BNFreeTransformTypeList(list, count);
+	return result;
+}
+
+
+BNTransformType Transform::GetType() const
+{
+	return BNGetTransformType(m_xform);
+}
+
+
+string Transform::GetName() const
+{
+	char* name = BNGetTransformName(m_xform);
+	string result = name;
+	BNFreeString(name);
+	return result;
+}
+
+
+string Transform::GetLongName() const
+{
+	char* name = BNGetTransformLongName(m_xform);
+	string result = name;
+	BNFreeString(name);
+	return result;
+}
+
+
+vector<TransformParameter> Transform::GetParameters() const
+{
+	return vector<TransformParameter>();
+}
+
+
+bool Transform::Decode(const DataBuffer& input, DataBuffer& output, const map<string, DataBuffer>& params)
+{
+	if (GetType() == InvertingTransform)
+		return Encode(input, output, params);
+	return false;
+}
+
+
+bool Transform::Encode(const DataBuffer&, DataBuffer&, const map<string, DataBuffer>&)
+{
+	return false;
+}
+
+
+CoreTransform::CoreTransform(BNTransform* xform): Transform(xform)
+{
+}
+
+
+vector<TransformParameter> CoreTransform::GetParameters() const
+{
+	size_t count;
+	BNTransformParameterInfo* list = BNGetTransformParameterList(m_xform, &count);
+
+	vector<TransformParameter> result;
+	for (size_t i = 0; i < count; i++)
+	{
+		TransformParameter param;
+		param.name = list[i].name;
+		param.longName = list[i].longName;
+		param.fixedLength = list[i].fixedLength;
+	}
+
+	BNFreeTransformParameterList(list, count);
+	return result;
+}
+
+
+bool CoreTransform::Decode(const DataBuffer& input, DataBuffer& output, const map<string, DataBuffer>& params)
+{
+	BNTransformParameter* list = new BNTransformParameter[params.size()];
+	size_t idx = 0;
+	for (auto i : params)
+	{
+		list[idx].name = i.first.c_str();
+		list[idx].value = i.second.GetBufferObject();
+	}
+
+	bool result = BNDecode(m_xform, input.GetBufferObject(), output.GetBufferObject(), list, idx);
+
+	delete[] list;
+	return result;
+}
+
+
+bool CoreTransform::Encode(const DataBuffer& input, DataBuffer& output, const map<string, DataBuffer>& params)
+{
+	BNTransformParameter* list = new BNTransformParameter[params.size()];
+	size_t idx = 0;
+	for (auto i : params)
+	{
+		list[idx].name = i.first.c_str();
+		list[idx].value = i.second.GetBufferObject();
+	}
+
+	bool result = BNEncode(m_xform, input.GetBufferObject(), output.GetBufferObject(), list, idx);
+
+	delete[] list;
+	return result;
+}
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 7c4aeb35..590202a6 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -6,6 +6,7 @@
 #include <stddef.h>
 #include <string>
 #include <vector>
+#include <map>
 #include <exception>
 #include "binaryninjacore.h"
 
@@ -422,6 +423,8 @@ namespace BinaryNinja
 		void SetEndianness(BNEndianness endian);
 
 		void Read(void* dest, size_t len);
+		DataBuffer Read(size_t len);
+		std::string ReadString(size_t len);
 		uint8_t Read8();
 		uint16_t Read16();
 		uint32_t Read32();
@@ -434,6 +437,8 @@ namespace BinaryNinja
 		uint64_t ReadBE64();
 
 		bool TryRead(void* dest, size_t len);
+		bool TryRead(DataBuffer& dest, size_t len);
+		bool TryReadString(std::string& dest, size_t len);
 		bool TryRead8(uint8_t& result);
 		bool TryRead16(uint16_t& result);
 		bool TryRead32(uint32_t& result);
@@ -472,6 +477,8 @@ namespace BinaryNinja
 		void SetEndianness(BNEndianness endian);
 
 		void Write(const void* src, size_t len);
+		void Write(const DataBuffer& buf);
+		void Write(const std::string& str);
 		void Write8(uint8_t val);
 		void Write16(uint16_t val);
 		void Write32(uint32_t val);
@@ -484,6 +491,8 @@ namespace BinaryNinja
 		void WriteBE64(uint64_t val);
 
 		bool TryWrite(const void* src, size_t len);
+		bool TryWrite(const DataBuffer& buf);
+		bool TryWrite(const std::string& str);
 		bool TryWrite8(uint8_t val);
 		bool TryWrite16(uint16_t val);
 		bool TryWrite32(uint32_t val);
@@ -499,4 +508,58 @@ namespace BinaryNinja
 		void Seek(uint64_t offset);
 		void SeekRelative(int64_t offset);
 	};
+
+	struct TransformParameter
+	{
+		std::string name, longName;
+		size_t fixedLength; // Variable length if zero
+	};
+
+	class Transform: public RefCountObject
+	{
+	protected:
+		BNTransform* m_xform;
+		BNTransformType m_typeForRegister;
+		std::string m_nameForRegister, m_longNameForRegister;
+
+		Transform(BNTransform* xform);
+
+		static BNTransformParameterInfo* GetParametersCallback(void* ctxt, size_t* count);
+		static void FreeParametersCallback(BNTransformParameterInfo* params, size_t count);
+		static bool DecodeCallback(void* ctxt, BNDataBuffer* input, BNDataBuffer* output, BNTransformParameter* params, size_t paramCount);
+		static bool EncodeCallback(void* ctxt, BNDataBuffer* input, BNDataBuffer* output, BNTransformParameter* params, size_t paramCount);
+
+		static std::vector<TransformParameter> EncryptionKeyParameters(size_t fixedKeyLength = 0);
+		static std::vector<TransformParameter> EncryptionKeyAndIVParameters(size_t fixedKeyLength = 0, size_t fixedIVLength = 0);
+
+	public:
+		Transform(BNTransformType type, const std::string& name, const std::string& longName);
+
+		static void Register(Transform* xform);
+		static Ref<Transform> GetByName(const std::string& name);
+		static std::vector<Ref<Transform>> GetTransformTypes();
+
+		BNTransformType GetType() const;
+		std::string GetName() const;
+		std::string GetLongName() const;
+
+		virtual std::vector<TransformParameter> GetParameters() const;
+
+		virtual bool Decode(const DataBuffer& input, DataBuffer& output, const std::map<std::string, DataBuffer>& params =
+			std::map<std::string, DataBuffer>());
+		virtual bool Encode(const DataBuffer& input, DataBuffer& output, const std::map<std::string, DataBuffer>& params =
+			std::map<std::string, DataBuffer>());
+	};
+
+	class CoreTransform: public Transform
+	{
+	public:
+		CoreTransform(BNTransform* xform);
+		virtual std::vector<TransformParameter> GetParameters() const override;
+
+		virtual bool Decode(const DataBuffer& input, DataBuffer& output, const std::map<std::string, DataBuffer>& params =
+			std::map<std::string, DataBuffer>()) override;
+		virtual bool Encode(const DataBuffer& input, DataBuffer& output, const std::map<std::string, DataBuffer>& params =
+			std::map<std::string, DataBuffer>()) override;
+	};
 }