text: #include "VorbisEncrypter.h"
  
  #include <time.h>
  #include "keys.h"
  #include "OggMap.h"
  
  typedef struct {
  	uint32_t a;
  	uint32_t b;
  } MapEntry;
  
  void VorbisEncrypter::GenerateIv(uint8_t* header_ptr) {
  	// Generate IV
  	srand(time(NULL));
  	initial_counter = (aes_ctr_128*)header_ptr;
  	// For convenience leave the low 4 IV bytes as zero
  	for (int i = 0; i < 4; i++) {
  		*header_ptr++ = 0;
  	}
  	// Generate the remaining 12 IV bytes with rand
  	for (int i = 0; i < 12; i++) {
  		*header_ptr++ = rand() & 0xFF;
  	}
  }
  
  VorbisEncrypter::VorbisEncrypter(void* datasource, ov_callbacks cbStruct) 
  	: file_ref(datasource), cb_struct(cbStruct) {
  	cb_struct.seek_func(file_ref, 0, SEEK_END);
  	uint32_t total_length = cbStruct.tell_func(file_ref);
  	cb_struct.seek_func(file_ref, 0, SEEK_SET);
  	struct {
  		int version;
  		int offset;
  	} original_file_header;
  	if (cbStruct.read_func(&original_file_header, sizeof(original_file_header), 1, file_ref) != 1)
  		throw std::exception("Unable to read mogg header.");
  	if (original_file_header.version != 0xA)
  		throw std::exception("Source mogg must be version 10/0xA (unencrypted).");
  
  	struct {
  		uint32_t version;
  		uint32_t chunk_size;
  		uint32_t num_entries;
  	} OggMapHdr;
  
  	if (cbStruct.read_func(&OggMapHdr, sizeof(OggMapHdr), 1, file_ref) != 1)
  		throw std::exception("Unable to read OggMap header.");
  	size_t header_size = 16 /* IV */
  		+ sizeof(original_file_header)
  		+ sizeof(OggMapHdr)
  		+ (OggMapHdr.num_entries * sizeof(MapEntry));
  	hmx_header.resize(header_size);
  
  	auto* header_ptr = hmx_header.data();
  	auto new_header = original_file_header;
  	new_header.version = 0xB;
  	new_header.offset = header_size;
  	// Copy Mogg header
  	memcpy(header_ptr, &new_header, sizeof(new_header));
  	header_ptr += sizeof(new_header);
  	// Copy OggMap header
  	memcpy(header_ptr, &OggMapHdr, sizeof(OggMapHdr));
  	header_ptr += sizeof(OggMapHdr);
  	// Copy OggMap data
  	cb_struct.read_func(header_ptr, sizeof(MapEntry), OggMapHdr.num_entries, file_ref);
  	header_ptr += sizeof(MapEntry) * OggMapHdr.num_entries;
  	GenerateIv(header_ptr);
  	encrypted_length = total_length - original_file_header.offset + hmx_header.size();
  	source_ogg_offset = original_file_header.offset;
  }
  
  
  VorbisEncrypter::VorbisEncrypter(void* datasource, int oggMapType, ov_callbacks cbStruct)
  	: file_ref(datasource), cb_struct(cbStruct)  {
  	cb_struct.seek_func(file_ref, 0, SEEK_END);
  	uint32_t total_length = cbStruct.tell_func(file_ref);
  	cb_struct.seek_func(file_ref, 0, SEEK_SET);
  
  	auto result = OggMap::Create(datasource, cbStruct);
  	if (std::holds_alternative<std::string>(result)){
  		throw std::exception(std::get<std::string>(result).c_str());
  	}
  	auto& map = std::get<OggMap>(result);
  	auto mapData = map.Serialize();
  	// 4 byte version, 4 byte offset, map, 16 byte IV
  	hmx_header.resize(8 + mapData.size() + 16);
  	auto *hdr_as_ints = (int32_t*)hmx_header.data();
  	hdr_as_ints[0] = 0xB;
  	hdr_as_ints[1] = (int32_t)hmx_header.size();
  	memcpy(&hdr_as_ints[2], mapData.data(), mapData.size());
  	GenerateIv(hmx_header.data() + hmx_header.size() - 16);
  
  	source_ogg_offset = 0;
  	encrypted_length = total_length + hmx_header.size();
  }
  
  VorbisEncrypter::~VorbisEncrypter() {
  	cb_struct.close_func(file_ref);
  }
  
  size_t VorbisEncrypter::ReadRaw(void* buf, size_t elementSize, size_t elements)
  {
  	size_t count = elementSize * elements;
  	size_t bytesRead = 0;
  	size_t offset = 0;
  	uint8_t* buffer = (uint8_t*)buf;
  
  	if (position < hmx_header.size()) {
  		while (position < hmx_header.size() && offset < count) {
  			buffer[offset++] = hmx_header[position++];
  			bytesRead++;
  		}
  	}
  	count -= offset;
  
  	if (position + count > encrypted_length) {
  		count = encrypted_length - position;
  	}
  	if (count == 0) return 0;
  
  	int32_t size_diff = (int32_t)source_ogg_offset - hmx_header.size();
  	cb_struct.seek_func(file_ref, position + size_diff, SEEK_SET);
  	size_t actualRead = cb_struct.read_func(buffer + offset, 1, count, file_ref);
  
  	bytesRead += actualRead;
  	position += actualRead;
  
  	EncryptBytes(buffer, offset, actualRead);
  	return bytesRead / elementSize;
  }
  
  
  
  // Fixes the counter
  void VorbisEncrypter::FixCounter(size_t decryptedPos) {
  	counter = *initial_counter;
  	uint64_t low = counter.qwords[0];
  	counter.qwords[0] += (decryptedPos >> 4);
  	if (counter.qwords[0] < low) {
  		counter.qwords[1]++;
  	}
  }
  
  // Increments a 128-bit counter using 64-bit word size
  inline void IncrementCounter(aes_ctr_128* ctr) {
  	for (int ptr = 0; ptr < 2; ptr++) {
  		ctr->qwords[ptr]++;
  		if (ctr->qwords[ptr] != 0)
  			break;
  	}
  }
  
  /**
   * buffer: buffer to write into
   * offset: offset into buffer to start writing
   * count: number of bytes to write into buffer
   *
   * This relies on the internal state, specifically that oggPos is set to the
   * last READ location. So if oggPos is 30 and count is 10 it assumes that bytes from 20 to 30 are being decrypted.
   */
  void VorbisEncrypter::EncryptBytes(uint8_t* buffer, size_t offset, size_t count)
  {
  	aes_ctr_128 cryptedCounter;
  	size_t decryptedPos = position - count - hmx_header.size();
  	int counterLoc = decryptedPos % 16;
  
  	FixCounter(decryptedPos);
  	AES128_ECB_encrypt(counter.bytes, ctrKey0B, cryptedCounter.bytes);
  	for (int i = 0; i < count; i++)
  	{
  		if (decryptedPos != 0 && decryptedPos % 16 == 0 && i != 0)
  		{
  			FixCounter(decryptedPos);
  			counterLoc = 0;
  			AES128_ECB_encrypt(counter.bytes, ctrKey0B, cryptedCounter.bytes);
  		}
  		buffer[i + offset] ^= cryptedCounter.bytes[counterLoc]; //decrypt one byte
  		decryptedPos++;
  		counterLoc++;
  	}
  }