def test_delimiter(self):
ex = Extractor(self.TEST_STRUCTURE, delimiter='|')
self.assertEqual(
ex.extract('key1|key2'),
self.TEST_STRUCTURE['key1']['key2']
)
self.assertEqual(
ex.extract('key1|key3.5'),
self.TEST_STRUCTURE['key1']['key3.5']
)
# This should return the fallback as its the wrong delimiter
self.assertEqual(
ex.extract('key1.key2'),
None
)
def test_source_priority(self):
ex = Extractor(self.TEST_STRUCTURE, self.TEST_STRUCTURE_B)
# The last source should take priority over the others
self.assertEqual(
ex.extract('key1.key2'),
self.TEST_STRUCTURE_B['key1']['key2']
)
# It should fallback to next source if last doesnt have
self.assertEqual(
ex.extract('struct1_only'),
self.TEST_STRUCTURE['struct1_only']
)
self.assertEqual(
ex.extract('key1.key4'),
self.TEST_STRUCTURE_B['key1']['key4']
)
def setUp(self):
self.ex = Extractor(self.TEST_STRUCTURE)
class TestExtractor(TestCase):
TEST_STRUCTURE = {
'key1': {
'key2': 'foo',
'key21': ['a', 'b'],
'key3.5': 'bar'
},
'struct1_only': True
}
TEST_STRUCTURE_B = {
'key1': {
'key2': 'zip',
'key21': ['x', 'y', 'z'],
'key3.5': 'ski',
'key4': 'code'
}
}
def setUp(self):
self.ex = Extractor(self.TEST_STRUCTURE)
def test_get_from_source(self):
# Succesful extractions
self.assertEqual(
self.ex.extract(['key1']),
self.TEST_STRUCTURE['key1']
)
self.assertEqual(
self.ex.extract(['key1', 'key2']),
self.TEST_STRUCTURE['key1']['key2']
)
self.assertEqual(
self.ex.extract(['key1', 'key21', '0']),
self.TEST_STRUCTURE['key1']['key21'][0]
)
self.assertEqual(
self.ex.extract('key1.key21.0'),
self.TEST_STRUCTURE['key1']['key21'][0]
)
self.assertEqual(
self.ex.extract(['key1','key3.5']),
self.TEST_STRUCTURE['key1']['key3.5']
)
def test_fallback(self):
# Fallback to defaults
test_default = 'foooooo'
self.assertEqual(
self.ex.extract(['BAD_KEY']),
None
)
self.assertEqual(
self.ex.extract(['BAD_KEY'], default=test_default),
test_default
)
self.assertEqual(
self.ex.extract(['key1', 'key21', '100']),
None
)
self.assertEqual(
self.ex.extract('foo.bar.zip.bad_key'),
None
)
with self.assertRaises(KeyError):
self.ex.extract('BAD.KEY', default=KeyError)
with self.assertRaises(KeyError):
self.ex.extract('BAD.KEY', default=KeyError('yooohooo'))
def test_wrapper(self):
self.assertEqual(
extract(self.TEST_STRUCTURE,'key1.key2'),
self.TEST_STRUCTURE['key1']['key2']
)
def test_source_priority(self):
ex = Extractor(self.TEST_STRUCTURE, self.TEST_STRUCTURE_B)
# The last source should take priority over the others
self.assertEqual(
ex.extract('key1.key2'),
self.TEST_STRUCTURE_B['key1']['key2']
)
# It should fallback to next source if last doesnt have
self.assertEqual(
ex.extract('struct1_only'),
self.TEST_STRUCTURE['struct1_only']
)
self.assertEqual(
ex.extract('key1.key4'),
self.TEST_STRUCTURE_B['key1']['key4']
)
def test_delimiter(self):
ex = Extractor(self.TEST_STRUCTURE, delimiter='|')
self.assertEqual(
ex.extract('key1|key2'),
self.TEST_STRUCTURE['key1']['key2']
)
self.assertEqual(
ex.extract('key1|key3.5'),
self.TEST_STRUCTURE['key1']['key3.5']
)
# This should return the fallback as its the wrong delimiter
self.assertEqual(
ex.extract('key1.key2'),
None
)
def vault_cli(source):
parser = argparse.ArgumentParser(description='Create update')
parser.add_argument('action', nargs='+')
parser.add_argument('-d', '--vault', required=True)
parser.add_argument('-k', '--key_db', default=None)
parser.add_argument('-u', '--user', default=getpass.getuser())
parser.add_argument('-i', '--private_key', default=None)
args = parser.parse_args(source)
action, action_args = args.action[0], args.action[1:]
vault = Vault(
args.vault,
key_db=args.key_db,
user=args.user,
private_key=args.private_key
)
if action == 'create':
vault.create()
elif action == 'view':
success = vault.load()
if not success:
print('Failed to decrypt')
else:
pprint(vault.data)
elif action == 'show_key':
success = vault.load()
if not success:
print('Failed to decrypt')
else:
print(vault.get_encryption_key())
elif action == 'edit':
success = vault.load()
if not success:
print('Failed to decrypt')
return False
tmp_file = args.vault + '.editing'
contents = json.dumps(vault.data, indent=4)
success = False
for i in range(5):
contents = user_edit(tmp_file, contents)
try:
new_data = json.loads(contents)
except Exception as e:
print('Edit Result Error:', str(e))
time.sleep(4)
else:
vault.data = new_data
success = True
break
if success:
vault.save()
else:
print('You failed too many times. Edits discarded.')
elif action == 'add_user':
success = vault.load()
if not success:
print('Failed to decrypt')
return False
user_name = action_args[0]
user_key = KeyPair(private_key_path=os.path.expanduser(action_args[1]))
vault.add_user(user_name, user_key)
vault.save()
elif action == 'rotate_key':
success = vault.load()
if not success:
print('Failed to decrypt')
return False
vault.change_key()
vault.save()
elif action == 'set':
vault.load()
pprint(vault.data)
elif action == 'get':
vault.load()
e = Extractor(vault.data)
value_path = action_args[0]
value = e.extract(value_path)
print('::{}::'.format(value_path))
pprint(value)
