Add flags

Author: eliben

2010/aes-encrypt-pycrypto/pycrypto_file.py

@@ -1,87 +1,98 @@
-import os
-import struct
-import random
-from eblib.utils import Timer
-from Crypto.Cipher import AES
-
-
-def encrypt_file(key, in_filename, out_filename=None, chunksize=64*1024):
-    """ Encrypts a file using AES (CBC mode) with the
-        given key.
-
-        key:
-            The encryption key - a bytes object that must be
-            either 16, 24 or 32 bytes long. Longer keys
-            are more secure.
-
-        in_filename:
-            Name of the input file
-
-        out_filename:
-            If None, '<in_filename>.enc' will be used.
-
-        chunksize:
-            Sets the size of the chunk which the function
-            uses to read and encrypt the file. Larger chunk
-            sizes can be faster for some files and machines.
-            chunksize must be divisible by 16.
-    """
-    if not out_filename:
-        out_filename = in_filename + '.enc'
-
-    iv = os.urandom(16)
-    encryptor = AES.new(key, AES.MODE_CBC, iv)
-    filesize = os.path.getsize(in_filename)
-
-    with open(in_filename, 'rb') as infile:
-        with open(out_filename, 'wb') as outfile:
-            outfile.write(struct.pack('<Q', filesize))
-            outfile.write(iv)
-
-            while True:
-                chunk = infile.read(chunksize)
-                if len(chunk) == 0:
-                    break
-                elif len(chunk) % 16 != 0:
-                    chunk += b' ' * (16 - len(chunk) % 16)
-
-                outfile.write(encryptor.encrypt(chunk))
-    return out_filename
-
-
-def decrypt_file(key, in_filename, out_filename=None, chunksize=24*1024):
-    """ Decrypts a file using AES (CBC mode) with the
-        given key. Parameters are similar to encrypt_file,
-        with one difference: out_filename, if not supplied
-        will be in_filename without its last extension
-        (i.e. if in_filename is 'aaa.zip.enc' then
-        out_filename will be 'aaa.zip')
-    """
-    if not out_filename:
-        out_filename = os.path.splitext(in_filename)[0]
-
-    with open(in_filename, 'rb') as infile:
-        origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
-        iv = infile.read(16)
-        decryptor = AES.new(key, AES.MODE_CBC, iv)
-
-        with open(out_filename, 'wb') as outfile:
-            while True:
-                chunk = infile.read(chunksize)
-                if len(chunk) == 0:
-                    break
-                outfile.write(decryptor.decrypt(chunk))
-
-            outfile.truncate(origsize)
-    return out_filename
-
-
-if __name__ == "__main__":
-    pwd = b'1' * 32
-    filename = '/tmp/file.my.enc'
-
-    #with Timer('encrypt'):
-        #enc_filename = encrypt_file(pwd, filename)
-
-    with Timer('decrypt'):
-        decrypt_file(pwd, filename, out_filename=filename+'.dec')
+import argparse
+import os
+import struct
+import random
+from eblib.utils import Timer
+from Crypto.Cipher import AES
+
+
+def encrypt_file(key, in_filename, out_filename=None, chunksize=64*1024):
+    """ Encrypts a file using AES (CBC mode) with the
+        given key.
+
+        key:
+            The encryption key - a bytes object that must be
+            either 16, 24 or 32 bytes long. Longer keys
+            are more secure.
+
+        in_filename:
+            Name of the input file
+
+        out_filename:
+            If None, '<in_filename>.enc' will be used.
+
+        chunksize:
+            Sets the size of the chunk which the function
+            uses to read and encrypt the file. Larger chunk
+            sizes can be faster for some files and machines.
+            chunksize must be divisible by 16.
+    """
+    if not out_filename:
+        out_filename = in_filename + '.enc'
+
+    iv = os.urandom(16)
+    encryptor = AES.new(key, AES.MODE_CBC, iv)
+    filesize = os.path.getsize(in_filename)
+
+    with open(in_filename, 'rb') as infile:
+        with open(out_filename, 'wb') as outfile:
+            outfile.write(struct.pack('<Q', filesize))
+            outfile.write(iv)
+
+            while True:
+                chunk = infile.read(chunksize)
+                if len(chunk) == 0:
+                    break
+                elif len(chunk) % 16 != 0:
+                    chunk += b' ' * (16 - len(chunk) % 16)
+
+                outfile.write(encryptor.encrypt(chunk))
+    return out_filename
+
+
+def decrypt_file(key, in_filename, out_filename=None, chunksize=24*1024):
+    """ Decrypts a file using AES (CBC mode) with the
+        given key. Parameters are similar to encrypt_file,
+        with one difference: out_filename, if not supplied
+        will be in_filename without its last extension
+        (i.e. if in_filename is 'aaa.zip.enc' then
+        out_filename will be 'aaa.zip')
+    """
+    if not out_filename:
+        out_filename = os.path.splitext(in_filename)[0]
+
+    with open(in_filename, 'rb') as infile:
+        origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
+        iv = infile.read(16)
+        decryptor = AES.new(key, AES.MODE_CBC, iv)
+
+        with open(out_filename, 'wb') as outfile:
+            while True:
+                chunk = infile.read(chunksize)
+                if len(chunk) == 0:
+                    break
+                outfile.write(decryptor.decrypt(chunk))
+
+            outfile.truncate(origsize)
+    return out_filename
+
+
+if __name__ == "__main__":
+    argparser = argparse.ArgumentParser(description='Encrypt/Decrypt AES')
+    argparser.add_argument('filename', nargs=1)
+    argparser.add_argument('-e', dest='encrypt', action='store_true')
+    argparser.add_argument('-d', dest='decrypt', action='store_true')
+    args = argparser.parse_args()
+    infile = args.filename[0]
+
+    key = b'1' * 32
+
+    if args.encrypt:
+        ofname = encrypt_file(key, infile, out_filename=infile+'.enc')
+        print('Encrypted to', ofname)
+    elif args.decrypt:
+        ofname = decrypt_file(key, infile, out_filename=infile+'.dec')
+        print('Decrypted to', ofname)
+    else:
+        argparser.print_help()
+        os.Exit(1)