This commit is contained in:
Luke Rogers 2013-10-01 13:38:00 +13:00
parent 72dce244b4
commit 46f571382c

View file

@ -23,36 +23,41 @@ def db_init(db):
def get_salt(bot): def get_salt(bot):
if not bot.config.get("random_salt", False): if not bot.config.get("random_salt", False):
bot.config["random_salt"] = hashlib.md5(os.urandom(16)).hexdigest() bot.config["random_salt"] = hashlib.md5(os.urandom(16)).hexdigest()
json.dump(bot.config, open('config', 'w'), sort_keys=True, indent=2) json.dump(bot.config, open('config', 'w'), sort_keys=True, indent=2)
return bot.config["random_salt"] return bot.config["random_salt"]
@hook.command @hook.command
def encrypt(inp, bot=None, db=None, notice=None): def encrypt(inp, bot=None, db=None, notice=None):
"""encrypt <pass> <string> -- Encrypts <string> with <pass>.""" """encrypt <pass> <string> -- Encrypts <string> with <pass>. (<string> can only be decrypted using this bot)"""
if not db_ready: if not db_ready:
db_init(db) db_init(db)
split = inp.split(" ") split = inp.split(" ")
# if there is only one argument, return the help message
if len(split) == 1: if len(split) == 1:
notice(encrypt.__doc__) notice(encrypt.__doc__)
return return
# generate the key from the password and salt
password = split[0] password = split[0]
salt = get_salt(bot) salt = get_salt(bot)
key = PBKDF2(password, salt) key = PBKDF2(password, salt)
# generate the IV and encode it to store in the database
iv = Random.new().read(AES.block_size); iv = Random.new().read(AES.block_size);
iv_encoded = base64.b64encode(iv) iv_encoded = base64.b64encode(iv)
# create the AES cipher and encrypt/encode the text with it
text = " ".join(split[1:]) text = " ".join(split[1:])
cipher = AES.new(key, AES.MODE_CBC, iv) cipher = AES.new(key, AES.MODE_CBC, iv)
encrypted = cipher.encrypt(pad(text)) encrypted = cipher.encrypt(pad(text))
encoded = base64.b64encode(encrypted) encoded = base64.b64encode(encrypted)
# store the encoded text and IV in the DB for decoding later
db.execute("insert or replace into encryption(encrypted, iv)" db.execute("insert or replace into encryption(encrypted, iv)"
"values(?,?)", (encoded, iv_encoded)) "values(?,?)", (encoded, iv_encoded))
db.commit() db.commit()
@ -62,25 +67,29 @@ def encrypt(inp, bot=None, db=None, notice=None):
@hook.command @hook.command
def decrypt(inp, bot=None, db=None, notice=None): def decrypt(inp, bot=None, db=None, notice=None):
"""decrypt <pass> <string> -- Decrypts <string> with <pass>.""" """decrypt <pass> <string> -- Decrypts <string> with <pass>. (can only decrypt strings encrypted on this bot)"""
if not db_ready: if not db_ready:
db_init(db) db_init(db)
split = inp.split(" ") split = inp.split(" ")
# if there is only one argument, return the help message
if len(split) == 1: if len(split) == 1:
notice(decrypt.__doc__) notice(decrypt.__doc__)
return return
# generate the key from the password and salt
password = split[0] password = split[0]
salt = get_salt(bot) salt = get_salt(bot)
key = PBKDF2(password, salt) key = PBKDF2(password, salt)
text = " ".join(split[1:]) text = " ".join(split[1:])
# get the encoded IV from the database and decode it
iv_encoded = db.execute("select iv from encryption where" iv_encoded = db.execute("select iv from encryption where"
" encrypted=?", (text,)).fetchone()[0] " encrypted=?", (text,)).fetchone()[0]
iv = base64.b64decode(iv_encoded) iv = base64.b64decode(iv_encoded)
# create AES cipher, decode text, decrypt text, and unpad it
cipher = AES.new(key, AES.MODE_CBC, iv) cipher = AES.new(key, AES.MODE_CBC, iv)
return unpad(cipher.decrypt(base64.b64decode(text))) return unpad(cipher.decrypt(base64.b64decode(text)))