def __init__(self, hsh, title, username, reason, description, url, performed, comment, authenticated, status): # type: (str, str, str, str, str, bool, str, bool, int) -> None ''' Creates a new Task for an alert that should go to `username` and is currently set to `status`. Args: title (str): The title of this task. username (str): The user who should be alerted from the Task. reason (str): The reason that the alert was fired. description (str): A description of the alert in question. url (str): A URL in which more information can be found about the alert itself, not the Task. performed (bool): Whether or not the user performed the action that caused this alert. comment (str): The user's comment on why the action occured. authenticated (bool): Whether 2FA has suceeded. status (enum): See `STATUS_LEVELS` from above. ''' self.title = title self.username = username self.reason = reason self.description = description self.url = url self.performed = performed self.comment = comment self.authenticated = authenticated self.status = status self._db_engine = DbEngine() self.hash = hsh
def __init__(self): # type: () -> None ''' Creates a new blacklist tied to a table named "blacklist". ''' # Load from table self._db_engine = DbEngine() names = self._db_engine.execute(config['queries']['blacklist_list']) # Break tuples into names self._blacklist = {name[0] for name in names}
class Tasker(object): ''' A simple class to retrieve tasks on which the bot should act upon. ''' def __init__(self): self._db_engine = DbEngine() def _get_tasks(self, level) -> List[Task]: # type: (int) -> List[Task] ''' Gets all tasks of a certain level. Args: level (int): One of StatusLevel Returns: List of SQLTasks. ''' alerts = self._db_engine.execute(config['queries']['get_alerts'], (level, )) return [Task(*alert) for alert in alerts] def get_new_tasks(self): # type: () -> List[Task] return self._get_tasks(StatusLevel.OPEN.value) def get_active_tasks(self): # type: () -> List[Task] return self._get_tasks(StatusLevel.INPROGRESS.value) def get_pending_tasks(self): # type: () -> List[Task] return self._get_tasks(StatusLevel.VERIFICATION.value)
def __update_ignored_list() -> None: # type: () -> None ''' Prunes the ignored table of old ignored alerts. ''' if config['queries'].get('update_ignored_list', False): DbEngine().execute(config['queries']['update_ignored_list'])
class Blacklist(object): def __init__(self): # type: () -> None ''' Creates a new blacklist tied to a table named "blacklist". ''' # Load from table self._db_engine = DbEngine() names = self._db_engine.execute(config['queries']['blacklist_list']) # Break tuples into names self._blacklist = {name[0] for name in names} def is_present(self, name): # type: (str) -> bool ''' Checks if a name is on the blacklist. Args: name (str): The name to check. ''' return name in self._blacklist def add(self, name): # type: (str) -> None ''' Adds a name to the blacklist. Args: name (str): The name to add to the blacklist. ''' self._blacklist.add(name) self._db_engine.execute(config['queries']['blacklist_add'], (name, )) def remove(self, name): # type: (str) -> None ''' Removes a name to the blacklist. Args: name (str): The name to remove from the blacklist. ''' self._blacklist.remove(name) self._db_engine.execute(config['queries']['blacklist_remove'], (name, ))
def get_ignored(username: str) -> Dict[str, str]: ''' Returns a dictionary of ignored alerts to reasons why the ignored are ignored. Args: username (str): The username of the user to retrieve ignored alerts for. Returns: Dict[str, str]: A mapping of ignored alert titles to reasons ''' __update_ignored_list() rows = DbEngine().execute(config['queries']['get_ignored'], (username, )) return {row[0]: row[1] for row in rows}
def create_new_alert(title, ldap, description, reason, url='N/A', key=None): # type: (str, str, str, str, str, str) -> None ''' Creates a new alert in the SQL DB with an optionally random hash. ''' # Generate random key if none provided if key is None: key = binascii.hexlify(os.urandom(32)) db_engine = DbEngine() # Insert that into the database as a new alert db_engine.execute(config['queries']['new_alert_alerts'], (key, ldap, title, description, reason, url)) db_engine.execute(config['queries']['new_alert_user_response'], (key,)) db_engine.execute(config['queries']['new_alert'], (key, StatusLevel.OPEN.value))
def ignore_task(username: str, title: str, reason: str, ttl: timedelta) -> None: ''' Adds a task with the given title to the ignore list for the given amount of time. Additionally adds an optional message to specify the reason that the alert was ignored. Args: username (str): The username of the user to ignore the given alert for. title (str): The title of the alert to ignore. ttl (Timedelta): The amount of time to ignore the alert for. msg (str): An optional string specifying why an alert was ignored ''' expiry_time = datetime.now(tz=pytz.utc) + ttl # NB: Non-standard MySQL specific query DbEngine().execute( config['queries']['ignore_task'], (username, title, reason, expiry_time.strftime('%Y-%m-%d %H:%M:%S')))
def main(): init() # init_sql() # Create components needed for SecurityBot duo_api = duo_client.Auth(ikey=config['duo']['ikey'], skey=config['duo']['skey'], host=config['duo']['endpoint']) duo_builder = lambda name: DuoAuth(duo_api, name) try: # Initialise DbEngine here DbEngine(config['database']) except KeyError: logging.error('No database configuration') raise chat = Slack(config['slack']) tasker = Tasker() sb = SecurityBot(chat, tasker, duo_builder, config['slack']['reporting_channel']) sb.run()
def __init__(self): self._db_engine = DbEngine()
class Task(object): def __init__(self, hsh, title, username, reason, description, url, performed, comment, authenticated, status): # type: (str, str, str, str, str, bool, str, bool, int) -> None ''' Creates a new Task for an alert that should go to `username` and is currently set to `status`. Args: title (str): The title of this task. username (str): The user who should be alerted from the Task. reason (str): The reason that the alert was fired. description (str): A description of the alert in question. url (str): A URL in which more information can be found about the alert itself, not the Task. performed (bool): Whether or not the user performed the action that caused this alert. comment (str): The user's comment on why the action occured. authenticated (bool): Whether 2FA has suceeded. status (enum): See `STATUS_LEVELS` from above. ''' self.title = title self.username = username self.reason = reason self.description = description self.url = url self.performed = performed self.comment = comment self.authenticated = authenticated self.status = status self._db_engine = DbEngine() self.hash = hsh def _set_status(self, status): # type: (int) -> None ''' Sets the status of a task in the DB. Args: status (int): The new status to use. ''' self._db_engine.execute(config['queries']['set_status'], (status, self.hash)) def _set_response(self): # type: () -> None ''' Updates the user response for this task. ''' self._db_engine.execute( config['queries']['set_response'], (self.comment, self.performed, self.authenticated, self.hash)) def set_open(self): self._set_status(StatusLevel.OPEN.value) def set_in_progress(self): self._set_status(StatusLevel.INPROGRESS.value) def set_verifying(self): self._set_status(StatusLevel.VERIFICATION.value) self._set_response()