diff --git a/plugins/encrypt.py b/plugins/encrypt.py new file mode 100644 index 0000000..95bdf44 --- /dev/null +++ b/plugins/encrypt.py @@ -0,0 +1,102 @@ +from util import hook +from Crypto import Random +from Crypto.Cipher import AES +from Crypto.Protocol.KDF import PBKDF2 + +import os +import base64 +import json +import hashlib + +# helper functions to pad and unpad a string to a specified block size +# +BS = AES.block_size +pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS) +unpad = lambda s : s[0:-ord(s[-1])] + +# helper functions to encrypt and encode a string with AES and base64 +encode_aes = lambda c, s: base64.b64encode(c.encrypt(pad(s))) +decode_aes = lambda c, s: unpad(c.decrypt(base64.b64decode(s))) + +db_ready = False + + +def db_init(db): + """check to see that our db has the the encryption table.""" + db.execute("create table if not exists encryption(encrypted, iv, " + "primary key(encrypted))") + db.commit() + db_ready = True + + +def get_salt(bot): + """generate an encryption salt if none exists, then returns the salt""" + if not bot.config.get("random_salt", False): + bot.config["random_salt"] = hashlib.md5(os.urandom(16)).hexdigest() + json.dump(bot.config, open('config', 'w'), sort_keys=True, indent=2) + return bot.config["random_salt"] + + +@hook.command +def encrypt(inp, bot=None, db=None, notice=None): + """encrypt -- Encrypts with . ( can only be decrypted using this bot)""" + if not db_ready: + db_init(db) + + split = inp.split(" ") + + # if there is only one argument, return the help message + if len(split) == 1: + notice(encrypt.__doc__) + return + + # generate the key from the password and salt + password = split[0] + salt = get_salt(bot) + key = PBKDF2(password, salt) + + # generate the IV and encode it to store in the database + iv = Random.new().read(AES.block_size); + iv_encoded = base64.b64encode(iv) + + # create the AES cipher and encrypt/encode the text with it + text = " ".join(split[1:]) + cipher = AES.new(key, AES.MODE_CBC, iv) + encoded = encode_aes(cipher, text) + + # store the encoded text and IV in the DB for decoding later + db.execute("insert or replace into encryption(encrypted, iv)" + "values(?,?)", (encoded, iv_encoded)) + db.commit() + + return encoded + + +@hook.command +def decrypt(inp, bot=None, db=None, notice=None): + """decrypt -- Decrypts with . (can only decrypt strings encrypted on this bot)""" + if not db_ready: + db_init(db) + + split = inp.split(" ") + + # if there is only one argument, return the help message + if len(split) == 1: + notice(decrypt.__doc__) + return + + # generate the key from the password and salt + password = split[0] + salt = get_salt(bot) + key = PBKDF2(password, salt) + + text = " ".join(split[1:]) + + # get the encoded IV from the database and decode it + iv_encoded = db.execute("select iv from encryption where" + " encrypted=?", (text,)).fetchone()[0] + iv = base64.b64decode(iv_encoded) + + # create AES cipher, decode text, decrypt text, and unpad it + cipher = AES.new(key, AES.MODE_CBC, iv) + return decode_aes(cipher, text) \ No newline at end of file