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 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])
class EventLog:
def __init__(self, db=None, logger=None):
self._cache = redis.Redis()
if db:
self._db = db
else:
self._db = Database(logger)
self._logger = logger
event_types = []
for each in self.list_event_types():
event_types.append(each[1])
missing_events = set(DEFAULT_EVENTS) - set(event_types)
for new_event in missing_events:
self._db.create_event_type(new_event)
def deserialize_user_object(self, event: Event):
if event.is_analyst():
cached_user_data = self._cache.get("analyst:{0}".format(
event.user_id))
if cached_user_data:
user_data = json.dumps(cached_user_data)
new_analyst_user = AnalystUser(
user_data["user_id"], user_data["email_address"],
user_data["session_token"], user_data["full_name"],
user_data["permissions"], user_data["view_event_type_ids"])
return new_analyst_user
else:
cached_user_data = self._cache.get("django:{0}".format(
event.user_id))
if cached_user_data:
user_data = json.dumps(cached_user_data)
new_django_user = DjangoUser(
user_data["user_id"], user_data["username"],
user_data["first_name"], user_data["last_name"],
user_data["email_address"],
datetime.datetime.strptime(user_data["created"],
"%Y-%m-%dT%H:%M:%S%z"),
datetime.datetime.strptime(user_data["last_logged_in"],
"%Y-%m-%dT%H:%M:%S%z"))
return new_django_user
def get_api_key_object(self, api_key: str) -> APIKey:
try:
node_data = self._db.validate_api_key(api_key)
node = APIKey(node_data[0], api_key, node_data[1], node_data[2],
node_data[3], node_data[4], node_data[5])
if node.suspended_event:
raise APIKeySuspended
if node.next_reset is None or node.next_reset < datetime.datetime.now(
):
next_reset = datetime.datetime.now() + datetime.timedelta(
hours=1)
node.next_reset = next_reset
node.events_posted = 0
self._db.reset_quota(node.node_id, next_reset)
if node.events_posted >= node.quota:
raise APIKeyRateLimited
return node
except DatabaseException:
raise APIKeyInvalid
def list_api_keys(self) -> list:
output = []
try:
all_api_keys = self._db.list_api_keys()
except DatabaseException:
raise EventLogException
for each_api_key in all_api_keys:
new_api_key = APIKey(each_api_key[0], each_api_key[1],
each_api_key[2], each_api_key[3],
each_api_key[4], each_api_key[5],
each_api_key[6])
output.append(new_api_key)
return output
def add_event_type(self, new_event_type):
self._cache.delete("event_types")
new_event_type_id = self._db.create_event_type(new_event_type)
if new_event_type_id > 0:
return new_event_type_id
raise EventLogException
@property
def event_type_ids_as_set(self) -> set:
try:
all_event_types = self.list_event_types()
all_event_type_ids = []
for each in all_event_types:
all_event_type_ids.append(each[0])
return set(all_event_type_ids)
except EventLogException:
raise EventLogException
def list_event_types(self) -> list:
event_types_cached = self._cache.get("event_types")
if event_types_cached:
event_types = json.loads(event_types_cached)
else:
event_types = self._db.list_event_types()
self._cache.set("event_types", json.dumps(event_types))
if len(event_types) == 1 and event_types[0][0] == -1:
raise EventLogException
return event_types
def get_event_count(self,
user_id=None,
event_type_id=None,
since=None,
until=None):
return self._db.get_event_count(user_id, event_type_id, since, until)
def log_event(self,
event_data: dict,
event_type_id: int,
user_id: int = None,
node_id: int = None):
event_id = self._db.log_event(event_data, event_type_id, user_id,
node_id)
if event_id:
event_type = None
for event in self.list_event_types():
if event[0] == event_type_id:
event_type = event[1]
break
if not event_type:
raise EventLogException
new_event = Event(event_id, event_type_id, event_type, event_data,
user_id, node_id,
datetime.datetime.now().isoformat())
self._cache.set("event_id:{0}".format(event_id), str(new_event))
return new_event
raise EventLogException
def get_event(self, event_id: int) -> Event:
cached_event_data = self._cache.get("event_id:{0}".format(event_id))
new_event = None
if cached_event_data:
cached_obj = json.loads(cached_event_data)
new_event = Event(
cached_obj["event_id"], cached_obj["event_type_id"],
cached_obj["event_type"], cached_obj["event_data"],
cached_obj["node_id"], cached_obj["user_id"],
datetime.datetime.strptime(cached_obj["created"],
"%Y-%m-%dT%H:%M:%S%z"))
else:
db_event_record = self._db.get_event_data(event_id)
if db_event_record:
# event_id, event_type_id, node_id, user_id, event_data, created, event_type.event_type
new_event = Event(db_event_record[0], db_event_record[1],
db_event_record[2],
json.loads(db_event_record[4]),
db_event_record[2], db_event_record[3],
db_event_record[5])
if new_event:
return new_event
else:
raise EventLogException
def retrieve_events(self,
user_id: int = None,
since: datetime = None,
until: datetime = None,
event_type_id: int = None,
node_id: int = None,
limit: int = 100):
try:
meta_data = self._db.retrieve_events(user_id, event_type_id,
node_id, since, until, limit)
# event_id, event_type_id, node_id, user_id, created
output = []
for row in meta_data:
event_id = row[0]
cache_key = "event_id:{0}".format(event_id)
event_data = self._cache.get(cache_key)
if event_data is None:
event_data = self._db.get_event_data(event_id)
new_event = Event(event_data[0],
event_data[1], event_data[6],
json.loads(event_data[4]), event_data[3],
event_data[2], event_data[5])
self._cache.set(cache_key, str(new_event))
output.append(json.loads(str(new_event)))
else:
output.append(json.loads(event_data))
return output
except DatabaseException:
raise EventLogException
def update_quota(self, node: APIKey):
try:
self._db.update_quota(node.node_id)
except DatabaseException:
if self._logger:
self._logger.error(
"Failed to update API rate quota for node id: {0}".format(
node.node_id))
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")
