Agent: Improve the speed of bit flipping code

- Remove a function call
- Use a generator
- Use a more efficient flip calculation (subtraction instead of xor)

Issue #2123

CHANGELOG.md

@@ -55,6 +55,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - `/api/auth` endpoint to `/api/authenticate`. #2105
 - `/api/registration` endpoint to `/api/register`. #2105
 - `/api/file-upload` endpoit to `/api/pba/upload`. #2154
+- Improved the speed of ransomware encryption by 2-3x. #2123
 
 ### Removed
 - VSFTPD exploiter. #1533

monkey/infection_monkey/utils/bit_manipulators.py

@@ -1,30 +1,21 @@
+from typing import Iterable
+
+
+def generate_flipped_bits(data: bytes) -> Iterable[int]:
+    """
+    Yield bytes with the bits flipped
+
+    :param data: The data whose bits to flip
+    """
+    for byte in data:
+        yield 255 - byte
+
+
 def flip_bits(data: bytes) -> bytes:
-    flipped_bits = bytearray(len(data))
+    """
+    Flip all bits in the given bytes
 
-    for i, byte in enumerate(data):
-        # TODO: The function call to flip_bits_in_single_byte() adds significant
-        #       overhead. While python is supposed to "inline" function calls
-        #       like this, I've yet to see any changes in runtime that indicate
-        #       this optimization is actually happening.
-        #
-        #       The value of breaking this into separate functions is the unit
-        #       test that tests all possible bytes (0-255). This gives us
-        #       confidence that our bit-flip operation is correct.
-        #
-        #       Remove the flip_bits_in_single_byte() function and rework the
-        #       unit tests so that we still have a high-degree of confidence
-        #       that this code is correct.
-        #
-        #       EDIT: I believe PyPy will attempt to inline functions
-        #       automatically. I don't know that CPython makes any such
-        #       optimizations.
-        flipped_bits[i] = flip_bits_in_single_byte(byte)
-
-    return bytes(flipped_bits)
-
-
-def flip_bits_in_single_byte(byte) -> int:
-    # TODO: The operation `255 - byte` appears to be 12% faster than 255 ^ byte.
-    #       Switch the operator and thoroughly test the ransomware payload to
-    #       ensure this doesn't introduce any defects.
-    return 255 ^ byte
+    :param data: The bytes whose bits to flip
+    :return: Bytes with the bits flipped
+    """
+    return bytes(generate_flipped_bits(data))

monkey/tests/unit_tests/infection_monkey/utils/test_bit_manipulators.py

@@ -1,9 +1,30 @@
 from infection_monkey.utils import bit_manipulators
+from infection_monkey.utils.bit_manipulators import flip_bits
 
 
-def test_flip_bits_in_single_byte():
-    for i in range(0, 256):
-        assert bit_manipulators.flip_bits_in_single_byte(i) == (255 - i)
+def test_flip_all_bits():
+    input_bytes = bytes(
+        b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f"
+        b"\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f"
+        b"\x20\x21\x22\x23\x24\x25\x26\x27\x28\x29\x2a\x2b\x2c\x2d\x2e\x2f"
+        b"\x30\x31\x32\x33\x34\x35\x36\x37\x38\x39\x3a\x3b\x3c\x3d\x3e\x3f"
+        b"\x40\x41\x42\x43\x44\x45\x46\x47\x48\x49\x4a\x4b\x4c\x4d\x4e\x4f"
+        b"\x50\x51\x52\x53\x54\x55\x56\x57\x58\x59\x5a\x5b\x5c\x5d\x5e\x5f"
+        b"\x60\x61\x62\x63\x64\x65\x66\x67\x68\x69\x6a\x6b\x6c\x6d\x6e\x6f"
+        b"\x70\x71\x72\x73\x74\x75\x76\x77\x78\x79\x7a\x7b\x7c\x7d\x7e\x7f"
+        b"\x80\x81\x82\x83\x84\x85\x86\x87\x88\x89\x8a\x8b\x8c\x8d\x8e\x8f"
+        b"\x90\x91\x92\x93\x94\x95\x96\x97\x98\x99\x9a\x9b\x9c\x9d\x9e\x9f"
+        b"\xa0\xa1\xa2\xa3\xa4\xa5\xa6\xa7\xa8\xa9\xaa\xab\xac\xad\xae\xaf"
+        b"\xb0\xb1\xb2\xb3\xb4\xb5\xb6\xb7\xb8\xb9\xba\xbb\xbc\xbd\xbe\xbf"
+        b"\xc0\xc1\xc2\xc3\xc4\xc5\xc6\xc7\xc8\xc9\xca\xcb\xcc\xcd\xce\xcf"
+        b"\xd0\xd1\xd2\xd3\xd4\xd5\xd6\xd7\xd8\xd9\xda\xdb\xdc\xdd\xde\xdf"
+        b"\xe0\xe1\xe2\xe3\xe4\xe5\xe6\xe7\xe8\xe9\xea\xeb\xec\xed\xee\xef"
+        b"\xf0\xf1\xf2\xf3\xf4\xf5\xf6\xf7\xf8\xf9\xfa\xfb\xfc\xfd\xfe\xff"
+    )
+
+    expected_bytes = bytes(reversed(input_bytes))
+
+    assert flip_bits(input_bytes) == expected_bytes
 
 
 def test_flip_bits():