encrypt, return bytes object instead of integer

.gitignore

@@ -4,3 +4,4 @@
 .swp.*
 
 *.pdf
+*.jpg

main.py

@@ -56,11 +56,17 @@ def css_encrypt(text, key):
     lfsr17 = lfsr((key[:16]+[1]), taps17)
     taps25 = [0, 3, 4, 12]
     lfsr25 = lfsr((key[16:]+[1]), taps25)
-    cipher_text = 0
+    cipher = 0
     carry = 0
-    bytes = text.to_bytes((text.bit_length() + 7) // 8, 'big')
 
-    for byte in bytes:
+    if isinstance(text, int):
+        Bytes = text.to_bytes((text.bit_length() + 7) // 8, 'big')
+    elif isinstance(text, bytes):
+        Bytes = text
+    else:
+        raise TypeError("input text should be bytes or integer")
+
+    for byte in Bytes:
         x_b2 = ""
         y_b2 = ""
         # Generate bytes from lfsr17 and lfsr25
@@ -74,17 +80,32 @@ def css_encrypt(text, key):
         carry = 1 if x + y > 255 else 0
 
         cipher_byte = z ^ byte
-        cipher_text = (cipher_text << 8) | cipher_byte
-    return cipher_text
+        # print(cipher_byte.to_bytes((cipher_byte.bit_length() + 7) // 8, 'big'))
+        cipher = (cipher << 8) | cipher_byte
+        cipher_bytes = b'\x00'*(len(Bytes) - len(cipher.to_bytes((cipher.bit_length() + 7) // 8, 'big'))) +cipher.to_bytes((cipher.bit_length() + 7) // 8, 'big') # padding
+    return cipher_bytes
 
 def test_encrypt():
     print("test encryption: text 0xffffffffff, key 0x0 ∈ {0, 1}^40")
-    cipher = css_encrypt(0xffffffffff, [0]*40)
+    cipher = int.from_bytes(css_encrypt(0xffffffffff, [0]*40), byteorder='big')
     print(f'cipher: {hex(cipher)}')
-    clear = css_encrypt(cipher, [0]*40)
+    clear = int.from_bytes(css_encrypt(cipher, [0]*40))
     print(f'decrypted: {hex(clear)}')
     print(f'original text and decrypted message are the same: {clear==0xffffffffff}')
 
+def gen_6_bytes():
+    # key = [randint(0, 1) for _ in range(40)]
+    # key = [1, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0]
+    key = [0]*40
+    text = b'\x00\x00\x00\x00\x00\x00'
+    print(key)
+    cipher = css_encrypt(text, key)
+    print(cipher)
+    print(css_encrypt(cipher, key))
+    return cipher
+
 test_lfsr17()
 print()
 test_encrypt()
+print()
+gen_6_bytes()