diff --git a/kaitai/kaitaistream.cpp b/kaitai/kaitaistream.cpp index 81d0d48..7cb1743 100644 --- a/kaitai/kaitaistream.cpp +++ b/kaitai/kaitaistream.cpp @@ -432,41 +432,108 @@ std::string kaitai::kstream::bytes_terminate(std::string src, char term, bool in // ======================================================================== std::string kaitai::kstream::process_xor_one(std::string data, uint8_t key) { + if (key == 0) + return data; + size_t len = data.length(); std::string result(len, ' '); - for (size_t i = 0; i < len; i++) + for (size_t i = 0; i < len; i++) { result[i] = data[i] ^ key; + } return result; } std::string kaitai::kstream::process_xor_many(std::string data, std::string key) { size_t len = data.length(); - size_t kl = key.length(); + if (len == 1) + return process_xor_one(data, key[0]); + std::string result(len, ' '); - size_t ki = 0; + size_t k = 0; + size_t keylen = key.length(); for (size_t i = 0; i < len; i++) { - result[i] = data[i] ^ key[ki]; - ki++; - if (ki >= kl) - ki = 0; + result[i] = data[i] ^ key[k]; + k++; + if (k == keylen) + k = 0; } return result; } -std::string kaitai::kstream::process_rotate_left(std::string data, int amount) { +uint8_t precomputedSingleRotations[8][256]; + +// NOTE: code based on StackOverflow answer at https://stackoverflow.com/a/34321324/2375119 +computeSingleRotations { + for (int amount = 1; amount < 8; amount++) { + int anti_amount = 8 - amount; + for (uint8_t i = 0; i < 256; i++) { + precomputedSingleRotations[amount][i] = (uint8_t)((i << amount) | (i >> anti_amount)); + } + } +} + +std::string kaitai::kstream::process_rotate_left(std::string data, int amount, int groupSize = 1) { + if (groupSize < 1) + throw std::runtime_error("process_rotate_left: groupSize must be at least 1"); + + amount = mod(amount, groupSize * 8); + if (amount == 0) + return data; + + int amount_bytes = amount / 8; size_t len = data.length(); std::string result(len, ' '); - for (size_t i = 0; i < len; i++) { - uint8_t bits = data[i]; - result[i] = (bits << amount) | (bits >> (8 - amount)); + if (groupSize == 1) { + uint8_t *translate = &precomputedSingleRotations[amount][0]; + + for (size_t i = 0; i < len; i++) { + result[i] = translate[data[i]]; + } + + return result; } - return result; + if (len % groupSize != 0) + throw std::runtime_error("process_rotate_left: data length must be a multiple of group size"); + + if (amount % 8 == 0) { + size_t indices[groupSize]; + for (size_t i = 0; i < groupSize; i++) { + indices[i] = (size_t)((i + amount_bytes) % groupSize); + } + + for (size_t i = 0; i < len; i += groupSize) { + for (size_t k = 0; k < groupSize; k++) { + result[i+k] = data[i + indices[k]]; + } + } + + return result; + } + + { + int amount1 = amount % 8; + int amount2 = 8 - amount1; + size_t indices1[groupSize]; + size_t indices2[groupSize]; + for (size_t i = 0; i < groupSize; i++) { + indices1[i] = (size_t)((i + amount_bytes) % groupSize); + indices2[i] = (size_t)((i + 1 + amount_bytes) % groupSize); + } + + for (size_t i = 0; i < len; i += groupSize) { + for (size_t k = 0; k < groupSize; k++) { + result[i+k] = (uint8_t)((data[i + indices1[k]] << amount1) | (data[i + indices2[k]] >> amount2)); + } + } + + return result; + } } #ifdef KS_ZLIB diff --git a/kaitai/kaitaistream.h b/kaitai/kaitaistream.h index 9592771..9cceaf7 100644 --- a/kaitai/kaitaistream.h +++ b/kaitai/kaitaistream.h @@ -172,6 +172,8 @@ class kstream { /** * Performs a XOR processing with given data, XORing every byte of input with a single * given value. + * WARNING: May return same byte array if key is zero. + * * @param data data to process * @param key value to XOR with * @return processed data @@ -182,6 +184,8 @@ class kstream { * Performs a XOR processing with given data, XORing every byte of input with a key * array, repeating key array many times, if necessary (i.e. if data array is longer * than key array). + * WARNING: May return same byte array if key is zero. + * * @param data data to process * @param key array of bytes to XOR with * @return processed data @@ -192,11 +196,14 @@ class kstream { * Performs a circular left rotation shift for a given buffer by a given amount of bits, * using groups of 1 bytes each time. Right circular rotation should be performed * using this procedure with corrected amount. + * WARNING: May return same byte array if amount is zero (modulo-wise). + * * @param data source data to process * @param amount number of bits to shift by + * @param groupSize number of bytes that make a group * @return copy of source array with requested shift applied */ - static std::string process_rotate_left(std::string data, int amount); + static std::string process_rotate_left(std::string data, int amount, int groupSize = 1); #ifdef KS_ZLIB /**