python file encryption

just ripped off a stackexchange post to make a quick file encrypt/decrypt tool.

from pdb import set_trace as dbg

from Crypto.Cipher import AES
from Crypto import Random
import base64
import hashlib
BS = 16 # couldn't get it to decrypt properly with 32
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS)
unpad = lambda s : s[0:-ord(s[-1])]
def passwdToKey(passwd):
return hashlib.sha256(passwd).digest()[:BS]

class AESCipher:
def __init__( self, key ):
self.key = key

def encrypt( self, raw ):
raw = pad(raw)
iv = Random.new().read( AES.block_size )
cipher = AES.new( self.key, AES.MODE_CBC, iv )
return base64.b64encode( iv + cipher.encrypt( raw ) )

def decrypt( self, enc ):
enc = base64.b64decode(enc)
iv = enc[:BS]
cipher = AES.new(self.key, AES.MODE_CBC, iv )
return unpad(cipher.decrypt( enc[BS:] ))

def executeAction(action, fileIn, fileOut, pw):
aesc = AESCipher(passwdToKey(pw))
if action == 'encrypt':
unencrypted = file(fileIn, 'rb').read()
encrypted = aesc.encrypt(unencrypted)
# verify encryption
assert aesc.decrypt(encrypted) == unencrypted, 'invalid encryption'
file(fileOut, 'wb').write(encrypted)
elif action == 'decrypt':
file(fileOut, 'wb').write(aesc.decrypt(file(fileIn, 'rb').read()))
else:
raise Exception('Unknown command: ' + action)
def main():
import sys
args = sys.argv[1:]
action, fileIn, fileOut, pw = args[-4:]
executeAction(action, fileIn, fileOut, pw)

if __name__ == '__main__':
main()