Permalink
Browse files

Do not require seeking to the end of the input to calculate padding, but

calculate padding after all bytes have been read
  • Loading branch information...
coleifer committed Sep 17, 2012
1 parent bf12073 commit 764b467076903130908cd9130438e86b4fe773e9
Showing with 10 additions and 13 deletions.
  1. +10 −13 beefish.py
View
@@ -15,14 +15,8 @@
from Crypto import Random
-def file_size(fh):
- fh.seek(0, 2)
- buflen = fh.tell()
- fh.seek(0)
- return buflen
-
-def _gen_padding(fh, block_size):
- pad_bytes = block_size - (file_size(fh) % block_size)
+def _gen_padding(file_size, block_size):
+ pad_bytes = block_size - (file_size % block_size)
padding = Random.get_random_bytes(pad_bytes - 1)
bflag = randrange(block_size - 2, 256 - block_size)
bflag -= bflag % block_size - pad_bytes
@@ -40,29 +34,31 @@ def get_cipher(key, iv):
def encrypt(in_buf, out_buf, key, chunk_size=4096):
iv = generate_iv(Blowfish.block_size)
cipher = get_cipher(key, iv)
- padding = _gen_padding(in_buf, cipher.block_size)
+ bytes_read = 0
wrote_padding = False
out_buf.write(iv)
while 1:
buffer = in_buf.read(chunk_size)
+ buffer_len = len(buffer)
+ bytes_read += buffer_len
if buffer:
- if len(buffer) < chunk_size:
- buffer += padding
+ if buffer_len < chunk_size:
+ buffer += _gen_padding(bytes_read, cipher.block_size)
wrote_padding = True
out_buf.write(cipher.encrypt(buffer))
else:
if not wrote_padding:
- out_buf.write(cipher.encrypt(padding))
+ out_buf.write(cipher.encrypt(_gen_padding(bytes_read, cipher.block_size)))
break
def decrypt(in_buf, out_buf, key, chunk_size=4096):
iv = in_buf.read(Blowfish.block_size)
cipher = get_cipher(key, iv)
decrypted = ''
-
+
while 1:
buffer = in_buf.read(chunk_size)
if buffer:
@@ -154,6 +150,7 @@ def test_stringio(self):
out_buf = StringIO()
dec_buf = StringIO()
in_buf.write(num_bytes * 'a')
+ in_buf.seek(0)
encrypt(in_buf, out_buf, self.key, i)
out_buf.seek(0)
decrypt(out_buf, dec_buf, self.key, i)

0 comments on commit 764b467

Please sign in to comment.