Merge commit '47eae6255989055440b45cfe667c9a6f8a92ce27' into develop
This commit is contained in:
commit
a6130ab780
102
plugins/encrypt.py
Normal file
102
plugins/encrypt.py
Normal file
|
@ -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
|
||||
# <http://stackoverflow.com/questions/12524994/encrypt-decrypt-using-pycrypto-aes-256>
|
||||
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 <pass> <string> -- Encrypts <string> with <pass>. (<string> 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 <pass> <string> -- Decrypts <string> with <pass>. (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)
|
Reference in a new issue