class TestCreateAPIKey(unittest.TestCase):
def setUp(self):
self._db = Database(logger, True)
def test_add_key(self):
new_api_key = binascii.hexlify(os.urandom(8)).decode('utf-8').upper()
result = self._db.create_api_key(new_api_key)
self.assertGreater(result, 0)
def test_invalid_api_key(self):
result = self._db.create_api_key("invalid")
self.assertEqual(result, -1)
def test_custom_reset_interval(self):
new_api_key = binascii.hexlify(os.urandom(8)).decode('utf-8').upper()
result = self._db.create_api_key(new_api_key, next_reset_seconds=5000)
self.assertGreater(result, 0)
api_key_info = self._db.validate_api_key(new_api_key)
lower_bounds = 4995
higher_bounds = 5005
self.assertGreater(api_key_info[3],
datetime.now() + timedelta(seconds=lower_bounds))
self.assertLess(api_key_info[3],
datetime.now() + timedelta(seconds=higher_bounds))
self.assertEqual(result, api_key_info[0])
class TestListEventTypes(unittest.TestCase):
def setUp(self) -> None:
# my default development environment values
dev = Environment("localhost", "localhost", "pico", "password",
"picoevent", "picoevent_test_db", 1000, 3600,
"localhost", "localhost")
self._db = Database(env=dev)
self._event_log = EventLog(self._db)
self._event_type_ids = self._event_log.list_event_types()
def test_listUnitTypes(self):
new_api_key = binascii.hexlify(os.urandom(8)).decode('utf-8').upper()
result = self._db.create_api_key(new_api_key,
rate_limit_quota=15000,
next_reset_seconds=5000)
self.assertGreater(result, 0)
rest_client = PicoEventRestClient.PicoEventRestClient(
TEST_ENDPOINT, new_api_key)
retrieved_event_types = rest_client.get_event_types()
self.assertEqual(len(retrieved_event_types), len(self._event_type_ids))
for x in range(0, len(retrieved_event_types)):
event_type = self._event_type_ids[x]
retrieved_event_type = retrieved_event_types[x]
self.assertEqual(event_type[0], retrieved_event_type[0])
self.assertEqual(event_type[1], retrieved_event_type[1])
def admin_create_api_key():
with current_app.app_context():
environment = Environment(current_app)
session_token = request.form["token"]
db = Database(logger=current_app.logger, env=environment)
analyst_user = db.validate_session(session_token)
if analyst_user and analyst_user.has_permission("create-api-key"):
new_api_key = binascii.hexlify(os.urandom(8)).decode('utf-8').upper()
result = db.create_api_key(new_api_key)
if result == -1:
flash("Could not create API key.", category="error")
return redirect(url_for("admin.home", session_token=session_token))
else:
abort(403)
class TestListAPIKeys(unittest.TestCase):
def setUp(self):
self._db = Database(logger, True)
self._target_api_keys = random.randint(5, 100)
self._api_keys = []
for x in range(0, self._target_api_keys):
new_api_key = binascii.hexlify(
os.urandom(8)).decode('utf-8').upper()
result = self._db.create_api_key(new_api_key)
self.assertGreater(result, 0)
self._api_keys.append(new_api_key)
def test_list_api_keys(self):
all_api_keys = self._db.list_api_keys()
self.assertEqual(self._target_api_keys, len(all_api_keys))
for x in range(0, self._target_api_keys):
self.assertEqual(all_api_keys[x][1], self._api_keys[x])
class TestRateLimitQuota(unittest.TestCase):
TEST_LIMIT = 500
def setUp(self):
self._db = Database(logger, True)
new_api_key = binascii.hexlify(os.urandom(8)).decode('utf-8').upper()
node_id = self._db.create_api_key(new_api_key)
self.assertGreater(node_id, 0)
self.api_key = new_api_key
self.node_id = node_id
def test_rateLimit(self):
for x in range(0, self.TEST_LIMIT):
self.assertTrue(self._db.update_quota(self.node_id))
node_info = self._db.validate_api_key(self.api_key)
self.assertEqual(node_info[0], self.node_id)
self.assertEqual(node_info[4], self.TEST_LIMIT)
next_reset = self._db.next_quota_reset
self.assertTrue(self._db.reset_quota(self.node_id, next_reset))
node_info = self._db.validate_api_key(self.api_key)
self.assertEqual(node_info[0], self.node_id)
self.assertEqual(0, node_info[4])
def display_menu(self):
if self._state == "USER_ADMIN":
self._user_admin_mgr.main_menu()
print("Console Administration Main Menu\n")
try:
test_db = Database(self._logger, True, False, self._env)
del test_db
except DatabaseException:
print(
"Warning: Test database does not exist or is not accessible with supplied credentials."
)
print("\n")
x = 1
for choice in ConsoleMainMenuManager.MAIN_MENU:
print("{0}.) {1}".format(x, choice))
x += 1
choice = self.select_item(1, 3)
if choice == 1:
self.change_administrator_password()
elif choice == 2:
_db = Database(self._logger, False, False, self._env)
# node_id, api_key, created, quota, next_reset, events_posted
all_api_keys = _db.list_api_keys()
del _db
print(
"Node ID\tAPI Key\tCreated\tQuota\tNext Reset\tEvents Posted")
for each_key in all_api_keys:
print("{0}\t{1}\t{2}\t{3}\t{4}\t{5}\t{6}".format(
each_key[0], each_key[1], each_key[2].isoformat(),
each_key[3], each_key[4], each_key[5]))
elif choice == 3:
print(
"Would you like to set a custom quota reset interval for the new API key?"
)
custom_reset_seconds = None
if self.confirm():
day_seconds = 3600 * 24
print(
"Choose a quota reset interval in seconds between 1 and {0}"
.format(day_seconds))
custom_reset_seconds = self.select_item(1, day_seconds)
print("Would you like to set a custom quota for the new API key?")
custom_quota = None
if self.confirm():
print("Choose a custom quota between 100 and 50000")
custom_quota = self.select_item(100, 50000)
new_api_key = binascii.hexlify(
os.urandom(8)).decode('utf-8').upper()
_db = Database(self._logger, False, False, self._env)
node_id = _db.create_api_key(new_api_key, custom_quota,
custom_reset_seconds)
if node_id > 0:
api_key_info = _db.validate_api_key(new_api_key)
print(
"New API Key {0} generated with a quota of {1}, next reset {1}."
.format(api_key_info[1], api_key_info[4], api_key_info[5]))
else:
print(
"Could not create a new API key, check log file for more info."
)
del _db
elif choice == 4:
print(
"Reset quota (you will need the node id from the API key list)"
)
_db = Database(self._logger, False, False, self._env)
all_node_ids = set()
for each_api_key in _db.list_api_keys:
all_node_ids.add(each_api_key[0])
repeat = True
node_id = None
while repeat:
user_input = int(input("Node ID: "))
if user_input in all_node_ids:
node_id = user_input
break
print(
"Node ID {0} not in the list of active API keys. Try again?"
)
repeat = self.confirm()
if node_id:
next_reset = _db.reset_quota
print(
"Would you like to set a custom reset interval (default: 3600 seconds)"
)
if self.confirm():
day_seconds = 3600 * 24
print(
"Choose a quota reset interval in seconds between 1 and {0}"
.format(day_seconds))
next_reset = datetime.now() + timedelta(
seconds=self.select_item(1, day_seconds))
_db.reset_quota(node_id, next_reset)
elif choice == 5:
self.user_admin_mode(self._user_id, self._session_token)
else:
print("Bye")
sys.exit(0)
input("Press Enter to continue")
