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)