From 0bf405d9613be5f86037b59b28fc9c29bcdbc825 Mon Sep 17 00:00:00 2001
From: mwiegand <mwiegand@seibert-media.net>
Date: Sun, 20 Jun 2021 20:22:47 +0200
Subject: [PATCH] wip

---
 libs/derive_string.py | 60 ++++++++++++++++++++++++++++---------------
 requirements.txt      |  1 +
 2 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/libs/derive_string.py b/libs/derive_string.py
index d4693f1..b59771e 100644
--- a/libs/derive_string.py
+++ b/libs/derive_string.py
@@ -3,29 +3,47 @@
 from hashlib import sha3_256
 from itertools import count
 from Crypto.Cipher import ChaCha20
-from math import ceil
+from math import floor, ceil
+
+def chacha_bits(input, bit_count):
+    zerobyte = (0).to_bytes(length=1, byteorder='big')
+    cipher = ChaCha20.new(key=sha3_256(input).digest(), nonce=zerobyte*8)
+    i = 0
+
+    while True:
+        start_bit = bit_count * i
+        start_byte = start_bit // 8
+        start_padding = start_bit % 8
+
+        end_bit = bit_count * i + bit_count
+        end_byte = end_bit // 8
+        end_padding = 8 - (end_bit % 8)
+
+        byte_count = end_byte - start_byte
+
+        cipher.seek(start_byte)
+        ciphertext = cipher.encrypt(zerobyte*byte_count)
+        shifted_ciphertext = int.from_bytes(ciphertext, byteorder='big') >> end_padding
+        
+        bit_mask = int('1'*bit_count, 2)
+        masked_ciphertext = shifted_ciphertext & bit_mask
+        
+        yield masked_ciphertext
+        i += 1
+
+
+def chacha_chracter(input, choices):
+    get_bits = chacha_bits(input, len(choices).bit_length())
+    
+    while True:
+        choice = next(get_bits)
+        if choice < len(choices):
+            yield choices[choice]
+
 
 def derive_string(input, length, choices):
-    result = ""
-    cipher = ChaCha20.new(key=sha3_256(input).digest())
-    
-    print(len(choices))
-    
-    for pow in count():
-        if 2**pow > len(choices):
-            break
-    
-    print(pow)
-
-    while len(result) < length:
-        seek = int.from_bytes(cipher.encrypt(b'0'*ceil(pow/8)), byteorder='little')
-        print(seek, len(choices))
-        if seek < len(choices):
-            result += choices[seek]
-        else:
-            continue
-
-    return result
+    get_character = chacha_chracter(input, choices)
+    return ''.join(next(get_character) for i in range(length))
 
 print(
     derive_string(b'12345', length=100, choices='abcde12345')
diff --git a/requirements.txt b/requirements.txt
index 93d2099..c0b176b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 bundlewrap>=4.8.0
+pycryptodome
 PyNaCl