def otp_test_custom_charset(charset):
test_lengths = [random.randint(1, 200), random.randint(200, 400), random.randint(400,2000)]
for length in test_lengths:
otp_tests(otp.securegen(length, charset=charset),
otp.securegen(length, charset=charset),
charset=charset
)
def cli_validate(encrypt_argv=['-m', otp.securegen(100)], decrypt_argv=['-m', '{}'], execute_args={}, include_key=True):
"""
Calls main() with encrypt and decrypt arguments.
decrypt_argv supports inserting encrypted_msg by setting -m to {}
"""
encrypt_args = new_argv([*encrypt_argv, '-oj', output_file])
result = execute(execute_args)
encrypted_msg = result['MESSAGE']
errors = result['ERRORS']
key = result['KEY']
e_args = array_to_dict(encrypt_args[1:])
msg = e_args.get('-m')
# If no msg and have a message file then load message file
if not msg and e_args.get('-f'):
args = {'encoding':execute_args.get('encoding', 'utf-8'), 'file_arg':'message_file', 'message_file': e_args.get('-f')}
msg = cli.CLI(args).get_message_from_file()
assert len(key) == len(encrypted_msg)
assert len(encrypted_msg) == len(msg)
assert len(errors) == 0
decrypt_argv = [ i.format(encrypted_msg) if '{}' else i for i in decrypt_argv ]
key_args = ['-k', key]
if not include_key:
key_args = []
decrypt_args = new_argv(['-d', *decrypt_argv, *key_args, '-oj', output_file])
decrypt_result = execute(execute_args)
assert msg == decrypt_result['MESSAGE']
assert len(decrypt_result['ERRORS']) == 0
assert decrypt_result['KEY'] == key
def test_encrypt_output_file(cleanup_test_file, cleanup_test_file_key, cleanup_output_file):
cli_validate(
decrypt_argv=['-f', test_file, '--charset', 'unicode', '-kf', test_file_key],
encrypt_argv=['-m', otp.securegen(100, charset=otp.utf_chars), '--charset', 'unicode', '-o', test_file, '-ok', test_file_key],
execute_args = {'encoding':'utf-16' },
include_key = False
)
def otp_tests(msg, key, **additional_args):
encrypted_msg, rkey = otp.otp(msg, key, **additional_args)
assert key == rkey
decrypted_msg, rkey = otp.otp(encrypted_msg, key, False, **additional_args)
assert key == rkey
assert decrypted_msg == msg
bad_decrypt, _ = otp.otp(encrypted_msg, otp.securegen(len(encrypted_msg)), False, **additional_args)
assert bad_decrypt != msg
def test_encrypt_input_file(cleanup_test_file, cleanup_test_file_key, cleanup_test_encrypted_file, cleanup_output_file):
for charset in util.charset_options.keys():
args = {'encoding':util.charset_get_encoding(charset), 'file_arg':'message_file', 'message_file':test_file}
cli.CLI(args).store_message_file(test_file, otp.securegen(100, charset=util.charset_options[charset]))
cli_validate(
decrypt_argv=['-f', test_encrypted_file, '-kf', test_file_key, '--charset', charset],
encrypt_argv=['-f', test_file, '-ok', test_file_key, '-o', test_encrypted_file, '--charset', charset],
execute_args = {'encoding': util.charset_get_encoding(charset)},
include_key = False
)
def test_encrypt(cleanup_output_file):
sys.argv = encrypt_string
cli_validate()
for option, charset in util.charset_options.items():
encoding = None
if option is 'unicode':
encoding = 'utf-16'
cli_validate(
decrypt_argv=['-m', '{}', '--charset', option],
encrypt_argv=['-m', otp.securegen(100, charset=charset), '--charset', option],
execute_args = {'encoding':encoding }
)
def validate_db_against_args(db, args):
items = db.items
assert len(items) == args.get('count')
item_keys = list(items.keys())
db.pop(item_keys[-1])
first_item_key = items[item_keys[0]]
assert len(items) == args.get('count') - 1
db.save()
db = ReferenceDB.load(db.file)
items = db.items
first_item_key_refresh = items[item_keys[0]]
otp_args = {}
securegen_args = {}
if args.get('charset'):
securegen_args['charset'] = args.get('charset')
otp_args['charset'] = args.get('charset')
otp_tests(otp.securegen(400, **securegen_args), first_item_key, **otp_args)
assert first_item_key == first_item_key_refresh
assert len(items) == args.get('count') - 1
first_key = item_keys[0]
assert len(first_key) == args.get('reference_key_length')
assert len(items.get(first_key)) == args.get('key_length')
ReferenceDB.delete(db.file)
assert os.path.exists(db.file) is False
def run(self):
if self.verbose:
print(self.args)
if self.generate or self.generate_list:
refdb_args = {}
self.refdb_args = refdb_args
if self.count:
refdb_args['count'] = int(self.count)
if self.key_length:
refdb_args['key_length'] = int(self.key_length)
refdb_args['charset'] = self.charset
if self.generate:
if self.reference_db:
refdb_args['file'] = self.reference_db
self.db = ReferenceDB.generate(**refdb_args)
print("Generated new reference database: {}".format(
self.db.file))
if self.generate_list:
if not self.reference_db_list:
raise ValueError(
'--generate_list requires directory for list in argument --reference_db_list'
)
refdb_args['directory'] = self.reference_db_list
self.dbl = ReferenceDBList.generate(files=int(
self.generate_list),
**refdb_args)
self.db = dbl.random_db()
if not self.message:
self.report_complete()
self.load_input_files(self.input_files)
if self.args['reference_db_list']:
self.dbl = ReferenceDBList(directory=self.reference_db_list)
if self.reference:
refparts = self.reference.split('-')
if len(refparts) == 2:
slug, self.reference = refparts
self.db = self.dbl.db_from_slug(slug)
else:
self.db = dbl.random_db()
self.eference_db = self.db.file
db_file = self.reference_db or DEFAULT_REFERENCE_DB_FILE
if not self.db:
if os.path.exists(db_file):
self.db = ReferenceDB.load(db_file)
if self.pop:
if not self.message:
if self.reference:
self.pop_reference_key()
self.report_complete()
else:
raise ValueError('--pop requires a reference key')
if not self.message:
raise ValueError(
'--message, --generate, or --pop is required. Try running --help'
)
noKey = not self.key and not self.reference
if self.decrypt and noKey:
raise ValueError('Decrypt requires either --key or --reference')
# If have reference key
if not self.key and not noKey:
self.key = self.db.items.get(self.reference)
# If no keys but just generated a database
generatedDB = noKey and self.generate
hasDbNoKey = not self.key and self.db
if generatedDB or hasDbNoKey:
self.reference, self.key = self.db.random_key()
# Ensure key is of sufficent length
if self.key and len(self.message) > len(
self.key) and not self.unsafe_key_length:
self.add_error(
"ERROR: KEY IS NOT OF SUFFICENT LENGTH TO ENCRYPT MESSAGE. Use --unsafe_key_length to disable"
)
self.report_complete(exit_code=1)
else:
if self.key:
if self.prefix_random_characters:
self.message = otp.securegen(
len(key) - len(self.message),
charset=self.charset) + self.message
otp_extra_args = {}
if self.charset == 'ascii' or self.charset == 'unicode':
otp_extra_args['charset'] = self.charset
mode = False if self.decrypt else True
self.message, self.key = otp.otp(msg=self.message,
key=self.key,
encrypt=mode,
**otp_extra_args)
# if isinstance(msg, bytes):
# msg = msg.decode('utf-16', 'surrogatepass')
reference_or_key = self.reference or self.key
reference_or_key = '{}-{}'.format(
db.slug,
reference_or_key) if self.reference_db_list else reference_or_key
self.key = reference_or_key
self.report_complete()
def encrypt_string():
yield sys.argv + ['-m', otp.securegen(100), '-oj', output_file]
os.remove(output_file)
def test_file_message_passing(cleanup_test_file):
messages = [otp.utf_chars[-1], otp.securegen(256, charset=otp.utf_chars)]
for message in messages:
args = {'encoding':'utf-16', 'file_arg':'message_file', 'message_file':test_file}
cli.CLI(args).store_message_file(test_file, message)
assert message == cli.CLI(args).get_message_from_file()
from cli import main
import cli
import sys
import pytest
import otp.otp as otp
import os
import random
from otp.util import array_to_dict, dict_key_value_swap
import otp.util as util
pkg = 'cli.py'
sys.argv = [pkg]
output_file = 'unittest-{}.json'.format(otp.securegen(8))
test_file = 'unittest-{}.otp'.format(otp.securegen(8))
test_file_key = 'unittest-{}.otp_key'.format(otp.securegen(8))
test_encrypted_file = 'unittest-{}.otp'.format(otp.securegen(8))
DATA_DIR = "{}/../data".format(os.path.dirname(os.path.abspath(__file__)))
@pytest.fixture
def encrypt_string():
yield sys.argv + ['-m', otp.securegen(100), '-oj', output_file]
os.remove(output_file)
@pytest.fixture
def cleanup_output_file():
yield
os.remove(output_file)
@pytest.fixture
def cleanup_test_file():
yield
os.remove(test_file)
@pytest.fixture
def generate_vc(**charset_args):
return otp.VigenereCipher(key=otp.securegen(random.randint(100, 300), **charset_args), msg=otp.securegen(random.randint(100,300), **charset_args), mode='encrypt', **charset_args)
def test_securegen(charsets):
assert len(otp.securegen(100)) == 100
# Test with a random two charsets
for charset in [ random.choice(charsets) for i in range(2) ]:
for i in otp.securegen(100, charset=charset):
assert i in charset
def test_simple():
otp_tests(otp.securegen(200), otp.securegen(200))
