timoreo/olm
1
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Handle non-base64 chars in pickle files

Browse source
This commit is contained in:
Richard van der Hoff 2016-09-22 15:09:40 +01:00
parent 047927d822
commit cfcee54a81
1 changed files with 18 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
python/olm/__main__.py
View file

@ -10,6 +10,11 @@ import yaml
from . import * from . import *
def read_base64_file(filename):
"""Read a base64 file, dropping any CR/LF characters"""
with open(filename, "rb") as f:
return f.read().translate(None, "\r\n")
def build_arg_parser(): def build_arg_parser():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--key", help="Account encryption key", default="") parser.add_argument("--key", help="Account encryption key", default="")
@ -37,8 +42,7 @@ def build_arg_parser():
def do_keys(args): def do_keys(args):
account = Account() account = Account()
with open(args.account_file, "rb") as f: account.unpickle(args.key, read_base64_file(args.account_file))
account.unpickle(args.key, f.read())
result = { result = {
"account_keys": account.identity_keys(), "account_keys": account.identity_keys(),
"one_time_keys": account.one_time_keys(), "one_time_keys": account.one_time_keys(),
@ -55,8 +59,7 @@ def build_arg_parser():
def do_id_key(args): def do_id_key(args):
account = Account() account = Account()
with open(args.account_file, "rb") as f: account.unpickle(args.key, read_base64_file(args.account_file))
account.unpickle(args.key, f.read())
print(account.identity_keys()['curve25519']) print(account.identity_keys()['curve25519'])
id_key = commands.add_parser("identity_key", help="Get the identity key for an account") id_key = commands.add_parser("identity_key", help="Get the identity key for an account")
@ -65,8 +68,7 @@ def build_arg_parser():
def do_one_time_key(args): def do_one_time_key(args):
account = Account() account = Account()
with open(args.account_file, "rb") as f: account.unpickle(args.key, read_base64_file(args.account_file))
account.unpickle(args.key, f.read())
keys = account.one_time_keys()['curve25519'].values() keys = account.one_time_keys()['curve25519'].values()
key_num = args.key_num key_num = args.key_num
if key_num < 1 or key_num > len(keys): if key_num < 1 or key_num > len(keys):
@ -93,8 +95,7 @@ def build_arg_parser():
def do_sign(args): def do_sign(args):
account = Account() account = Account()
with open(args.account_file, "rb") as f: account.unpickle(args.key, read_base64_file(args.account_file))
account.unpickle(args.key, f.read())
with open_in(args.message_file) as f: with open_in(args.message_file) as f:
message = f.read() message = f.read()
signature = account.sign(message) signature = account.sign(message)
@ -110,8 +111,7 @@ def build_arg_parser():
def do_generate_keys(args): def do_generate_keys(args):
account = Account() account = Account()
with open(args.account_file, "rb") as f: account.unpickle(args.key, read_base64_file(args.account_file))
account.unpickle(args.key, f.read())
account.generate_one_time_keys(args.count) account.generate_one_time_keys(args.count)
with open(args.account_file, "wb") as f: with open(args.account_file, "wb") as f:
f.write(account.pickle(args.key)) f.write(account.pickle(args.key))
@ -132,8 +132,7 @@ def build_arg_parser():
)) ))
sys.exit(1) sys.exit(1)
account = Account() account = Account()
with open(args.account_file, "rb") as f: account.unpickle(args.key, read_base64_file(args.account_file))
account.unpickle(args.key, f.read())
session = Session() session = Session()
session.create_outbound( session.create_outbound(
account, args.identity_key, args.one_time_key account, args.identity_key, args.one_time_key
@ -168,8 +167,7 @@ def build_arg_parser():
)) ))
sys.exit(1) sys.exit(1)
account = Account() account = Account()
with open(args.account_file, "rb") as f: account.unpickle(args.key, read_base64_file(args.account_file))
account.unpickle(args.key, f.read())
with open_in(args.message_file) as f: with open_in(args.message_file) as f:
message_type = f.read(8) message_type = f.read(8)
message = f.read() message = f.read()
@ -191,8 +189,7 @@ def build_arg_parser():
def do_session_id(args): def do_session_id(args):
session = Session() session = Session()
with open(args.session_file, "rb") as f: session.unpickle(args.key, read_base64_file(args.session_file))
session.unpickle(args.key, f.read())
sys.stdout.write(session.session_id() + "\n") sys.stdout.write(session.session_id() + "\n")
session_id.set_defaults(func=do_session_id) session_id.set_defaults(func=do_session_id)
@ -204,8 +201,7 @@ def build_arg_parser():
def do_encrypt(args): def do_encrypt(args):
session = Session() session = Session()
with open(args.session_file, "rb") as f: session.unpickle(args.key, read_base64_file(args.session_file))
session.unpickle(args.key, f.read())
with open_in(args.plaintext_file) as f: with open_in(args.plaintext_file) as f:
plaintext = f.read() plaintext = f.read()
message_type, message = session.encrypt(plaintext) message_type, message = session.encrypt(plaintext)
@ -224,8 +220,7 @@ def build_arg_parser():
def do_decrypt(args): def do_decrypt(args):
session = Session() session = Session()
with open(args.session_file, "rb") as f: session.unpickle(args.key, read_base64_file(args.session_file))
session.unpickle(args.key, f.read())
with open_in(args.message_file) as f: with open_in(args.message_file) as f:
message_type = f.read(8) message_type = f.read(8)
message = f.read() message = f.read()
@ -296,8 +291,7 @@ def do_outbound_group(args):
def do_group_encrypt(args): def do_group_encrypt(args):
session = OutboundGroupSession() session = OutboundGroupSession()
with open(args.session_file, "rb") as f: session.unpickle(args.key, read_base64_file(args.session_file))
session.unpickle(args.key, f.read())
plaintext = args.plaintext_file.read() plaintext = args.plaintext_file.read()
message = session.encrypt(plaintext) message = session.encrypt(plaintext)
with open(args.session_file, "wb") as f: with open(args.session_file, "wb") as f:
@ -306,8 +300,7 @@ def do_group_encrypt(args):
def do_group_credentials(args): def do_group_credentials(args):
session = OutboundGroupSession() session = OutboundGroupSession()
with open(args.session_file, "rb") as f: session.unpickle(args.key, read_base64_file(args.session_file))
session.unpickle(args.key, f.read())
result = { result = {
'message_index': session.message_index(), 'message_index': session.message_index(),
'session_key': session.session_key(), 'session_key': session.session_key(),
@ -333,8 +326,7 @@ def do_inbound_group(args):
def do_group_decrypt(args): def do_group_decrypt(args):
session = InboundGroupSession() session = InboundGroupSession()
with open(args.session_file, "rb") as f: session.unpickle(args.key, read_base64_file(args.session_file))
session.unpickle(args.key, f.read())
message = args.message_file.read() message = args.message_file.read()
plaintext = session.decrypt(message) plaintext = session.decrypt(message)
with open(args.session_file, "wb") as f: with open(args.session_file, "wb") as f: