diff --git a/kaitai/kaitaistream.cpp b/kaitai/kaitaistream.cpp index 81d0d48..7453d61 100644 --- a/kaitai/kaitaistream.cpp +++ b/kaitai/kaitaistream.cpp @@ -432,11 +432,15 @@ 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; } @@ -444,29 +448,80 @@ std::string kaitai::kstream::process_xor_one(std::string data, uint8_t key) { 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; 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[i % kl]; } return result; } -std::string kaitai::kstream::process_rotate_left(std::string data, int 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) { + int anti_amount = 8 - amount; + uint8_t translate[256]; + for (uint8_t i = 0; i < 256; i++) { + translate[i] = (uint8_t)((i << amount) | (i >> anti_amount)); + } + + 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[groupSize] indices; + 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[groupSize] indices1; + size_t[groupSize] indices2; + 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..024e52f 100644 --- a/kaitai/kaitaistream.h +++ b/kaitai/kaitaistream.h @@ -196,7 +196,7 @@ class kstream { * @param amount number of bits to shift by * @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 /**