Refactored entrypoint of library

afonsoc12 · Feb 12, 2022 · 69552ff · 69552ff
1 parent 82e8cd4
commit 69552ff
Show file tree

Hide file tree

Showing 3 changed files with 179 additions and 143 deletions.
diff --git a/intrusion_monitor/__init__.py b/intrusion_monitor/__init__.py
@@ -1,4 +1,153 @@
+import os
+import sys
+import logging
+from datetime import date
+from pathlib import Path
 
 from ._version import get_versions
+from .watchdog import Watchdog
+
+# Environment variables and if they are required
+ENVIRONMENT_VARS = {'TZ': False,
+                    'API_KEY': True,
+                    'INFLUXDB_HOST': False,
+                    'INFLUXDB_PORT': False,
+                    'INFLUXDB_DATABASE': False,
+                    'INFLUXDB_USER': False,
+                    'INFLUXDB_PASSWORD': False,
+                    'OPERATION_MODE': False,
+                    'SSH_LOG_PATH': False,
+                    'LOG_LEVEL': False}
+
 __version__ = get_versions()['version']
-del get_versions
+
+
+def _real_main():
+
+    # Setup logging
+    _logging_setup(copyright=True)
+
+    # Unset empty variables
+    _unset_empty_env(ENVIRONMENT_VARS)
+
+    # Check if required variables are present
+    _check_vars_exist(ENVIRONMENT_VARS)
+
+    # Select if working as a TCP socket (for rsyslog) or as a log watchdog (default)
+    OPERATION_MODE = os.getenv('OPERATION_MODE')
+    if not OPERATION_MODE:
+        logging.warning('OPERATION_MODE variable is not set. Defaulting to "watchdog"')
+        OPERATION_MODE = 'watchdog'
+    elif OPERATION_MODE.casefold() not in ('socket', 'watchdog'):
+        err = f'OPERATION_MODE={OPERATION_MODE} is not recognised and this cannot continue"'
+        logging.error(err)
+        raise EnvironmentError(err)
+    else:
+        logging.info(f'Using OPERATION_MODE={OPERATION_MODE}')
+
+    # Bootstrap intrusion-monitor from OPERATION_MODE
+    _bootstrap(OPERATION_MODE)
+
+
+def _bootstrap(operation_mode):
+    """Initialises intrusion-monitor in either `watchdog` or `socket` operation modes."""
+
+    if operation_mode == 'watchdog':
+
+        log_path = Path(os.getenv('SSH_LOG_PATH', '/watchdog/log/auth.log'))
+
+        # Check if file exists and can be read
+        if not log_path.exists():
+            err = f'No file was not found and this can\'t continue. Log path provided is: {log_path.absolute()}'
+            logging.critical(err)
+            return FileNotFoundError(err)
+        elif not os.access(log_path, os.R_OK):
+            err = f'The file cant be opened. Running: "sudo chmod o+r <Log file>" might solve this issue.'
+            logging.critical(err)
+            raise PermissionError(err)
+        else:
+            logging.info(f'Log file found at: {log_path.absolute()}')
+            with open(log_path, 'r') as f:
+                lines = f.readlines()
+                logging.debug('Here are the last 5 lines of the log file:\n\t{}'.format('\t'.join(lines[-5:])))
+
+            # Everything seems okay, starting watchdog
+            watchdog = Watchdog(log_path)
+            logging.debug(f'So far so good, starting log Watchdog...')
+            watchdog.start()
+
+    elif operation_mode == 'socket':
+        logging.critical(
+            f'This feature is not yet implemented and this can\'t continue. OPERATION_MODE is {operation_mode}')
+        raise NotImplementedError(f'The OPERATION_MODE={operation_mode} is not yet implemented.')
+        # server.start()
+    else:
+        logging.critical(
+            f'A critical problem occurred while trying to bootstrap from OPERATION_MODE and this can\'t continue. '
+            f'OPERATION_MODE is {operation_mode}')
+        raise EnvironmentError('A critical problem occurred while trying to bootstrap from OPERATION_MODE and this can\'t continue. ')
+
+
+def _unset_empty_env(vars):
+    """Unset empty environment variables."""
+
+    for v in vars:
+        var = os.getenv(v, None)
+        if not var and var is not None and len(var) == 0:
+            del os.environ[v]
+            logging.warning(f'Environment variable {v} is set but is empty. Unsetted...')
+
+
+def _logging_setup(copyright=True, version=True):
+    log_level = os.getenv('LOG_LEVEL', 'info')
+
+    if log_level.casefold() == 'debug':
+        log_level = logging.DEBUG
+    elif log_level.casefold() == 'info':
+        log_level = logging.INFO
+    else:
+        # Default
+        log_level = logging.INFO
+
+    logging.basicConfig(format='%(asctime)s [%(levelname)s]: %(message)s',
+                        datefmt='%Y-%m-%d %H:%M:%S',
+                        level=log_level)
+
+    # Print copyright
+    if copyright:
+        logging.info(f'Copyright {date.today().year} Afonso Costa')
+        logging.info('Licensed under the Apache License, Version 2.0 (the "License");')
+
+    # Print version
+    if version:
+        logging.info('Version: {}'.format(get_versions()['version']))
+
+
+def _check_vars_exist(vars):
+    """Checks if the required variables exist."""
+
+    vars_missing = []
+    for v in [v for v in vars if vars[v]]:
+        var = os.getenv(v, None)
+        if not var:
+            logging.error(f'Environment variable {v} is not set and its mandatory!')
+            vars_missing.append(v)
+
+    if vars_missing:
+        logging.critical(
+            'Some mandatory environment variables are not set and this can\'t continue. Env variables missing: {}'.format(
+                ', '.join(vars_missing)))
+        raise EnvironmentError('Some mandatory environment variables are not set. Env variables missing: {}'.format(
+            ', '.join(vars_missing)))
+
+
+def main():
+
+    try:
+        _real_main()
+    except KeyboardInterrupt:
+        logging.error('ERROR: Interrupted by user')
+        raise
+    except:
+        logging.critical('Fatal error occurred and intrusion-monitor cannot continue...')
+        raise
diff --git a/intrusion_monitor/__main__.py b/intrusion_monitor/__main__.py
@@ -1,142 +1,18 @@
-import os
-import sys
-import logging
-from datetime import date
-from pathlib import Path
-
-from ._version import get_versions
-from .watchdog import Watchdog
-
-# Environment variables and is required
-ENVIRONMENT_VARS = {'TZ': False,
-                    'API_KEY': True,
-                    'INFLUXDB_HOST': False,
-                    'INFLUXDB_PORT': False,
-                    'INFLUXDB_DATABASE': False,
-                    'INFLUXDB_USER': False,
-                    'INFLUXDB_PASSWORD': False,
-                    'OPERATION_MODE': False,
-                    'SSH_LOG_PATH': False,
-                    'LOG_LEVEL': False}
-
-
-def main():
-
-    # Setup logging
-    logging_setup()
-
-    logging.info(f'Copyright {date.today().year} Afonso Costa')
-    logging.info('Licensed under the Apache License, Version 2.0 (the "License");')
-    logging.info('Version: {}'.format(get_versions()['version']))
-
-    # Unset empty variables
-    unset_empty_env(ENVIRONMENT_VARS)
-
-    # Check if required variables are present
-    check_vars_exist(ENVIRONMENT_VARS)
-
-    # Select if working as a TCP socket (for rsyslog) or as a log watchdog (default)
-    OPERATION_MODE = os.getenv('OPERATION_MODE')
-    if not OPERATION_MODE:
-        logging.warning('OPERATION_MODE variable is not set. Defaulting to "watchdog"')
-        OPERATION_MODE = 'watchdog'
-    elif OPERATION_MODE.casefold() not in ('socket', 'watchdog'):
-        logging.warning(f'OPERATION_MODE={OPERATION_MODE} is not recognised. Defaulting to "watchdog"')
-        OPERATION_MODE = 'watchdog'
-    else:
-        logging.info(f'Using OPERATION_MODE={OPERATION_MODE}')
-
-    if OPERATION_MODE == 'watchdog':
-
-        log_path = Path(os.getenv('SSH_LOG_PATH', '/watchdog/log/auth.log'))
-
-        # Check if file exists and can be read
-        if not log_path.exists():
-            logging.critical(f'No file was found and this can\'t continue. Log path provided is: {log_path.absolute()}')
-            exit(127)
-        elif not os.access(log_path, os.R_OK):
-            logging.critical(f'The file cant be opened. Running: "sudo chmod o+r <Log file>" might solve this issue.')
-            exit(128)
-        else:
-            logging.info(f'Log file found at: {log_path.absolute()}')
-            with open(log_path, 'r') as f:
-                lines = f.readlines()
-                logging.debug('Here are the last 5 lines of the log file:\n\t{}'.format('\t'.join(lines[-5:])))
-
-            # Everything seems okay, starting watchdog
-            watchdog = Watchdog(log_path)
-            logging.debug(f'So far so good, starting log Watchdog...')
-            watchdog.start()
-    elif OPERATION_MODE == 'socket':
-        logging.critical(
-            f'This feature is not yet implemented and this can\'t continue. OPERATION_MODE is {OPERATION_MODE}')
-        raise NotImplementedError(f'The OPERATION_MODE={OPERATION_MODE} is not yet implemented.')
-        # server.start()
-    else:
-        logging.critical(
-            f'A critical problem occurred while handling OPERATION_MODE and this can\'t continue. OPERATION_MODE is {OPERATION_MODE}')
-        raise EnvironmentError('A problem occurred while handling OPERATION_MODE environment variable.')
+#!/usr/bin/env python3
 
-    exit(0)
+# Execute with
+# $ python intrusion_monitor/__main__.py
+# $ python -m intrusion_monitor
 
+import sys
 
-def unset_empty_env(vars):
-    """Unset empty environment variables."""
-
-    for v in vars:
-        var = os.getenv(v, None)
-        if not var and var is not None and len(var) == 0:
-            del os.environ[v]
-            logging.warning(f'Environment variable {v} is set but is empty. Unsetted...')
-
-
-def logging_setup():
-    log_level = os.getenv('LOG_LEVEL', 'info')
-
-    if log_level.casefold() == 'debug':
-        log_level = logging.DEBUG
-    elif log_level.casefold() == 'info':
-        log_level = logging.INFO
-    else:
-        # Default
-        log_level = logging.INFO
-
-    logging.basicConfig(format='%(asctime)s [%(levelname)s]: %(message)s',
-                        datefmt='%Y-%m-%d %H:%M:%S',
-                        level=log_level)
-
-
-def check_vars_exist(vars):
-    """Checks if the required variables exist."""
-
-    vars_missing = []
-    for v in [v for v in vars if vars[v]]:
-        var = os.getenv(v, None)
-        if not var:
-            logging.error(f'Environment variable {v} is not set and its mandatory!')
-            vars_missing.append(v)
-
-    if vars_missing:
-        logging.critical(
-            'Some mandatory environment variables are not set and this can\'t continue. Env variables missing: {}'.format(
-                ', '.join(vars_missing)))
-        raise EnvironmentError('Some mandatory environment variables are not set. Env variables missing: {}'.format(
-            ', '.join(vars_missing)))
-
-
-def exit(exit_code):
-    if exit_code == 0:
-        logging.info(f'Exiting with exit code {exit_code}')
-    else:
-        logging.error(f'Exiting with exit code {exit_code}')
-
-    sys.exit(exit_code)
+if __package__ is None and not hasattr(sys, 'frozen'):
+    # Direct call of __main__.py
+    import os.path
+    path = os.path.realpath(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.dirname(os.path.dirname(path)))
 
+import intrusion_monitor
 
 if __name__ == '__main__':
-
-    try:
-        main()
-    except Exception as e:
-        logging.critical('An exception occurred', exc_info=True)
-        exit(250)
+    intrusion_monitor.main()
diff --git a/intrusion_monitor/log_parser.py → intrusion_monitor/log_line.py b/intrusion_monitor/log_parser.py → intrusion_monitor/log_line.py
@@ -31,18 +31,20 @@ def _extract_log_parts(self):
         prefix, message = self.log_line.split(pid_grp[0]+': ')
 
         logging.debug(f'Log prefix extracted: {prefix}')
-        logging.debug(f'Log prefix extracted: {message}')
+        logging.debug(f'Log message extracted: {message}')
         return prefix, message
 
     def is_login_attempt(self):
-        """Check if the log line is considered a login attempt"""
+        """Check if the log line is considered a login attempt.
+
+           Returns a tuple with two values: line is login attempt and reason to consider as such.
+        """
         possibilities = ('Connection closed by', 'Failed password for', )
 
-        status = [False, None]
+        status = (False, None,)
         for p in possibilities:
             if self.message.startswith(p):
-                status[0] = True
-                status[1] = f'Matches {p}'
+                status = (True, p,)
                 break
 
         return status
@@ -195,5 +197,14 @@ def attempt_username(self):
         logging.debug(f'Username extracted: {username}')
         return username
 
-    def get_ip_info(self):
+    def get_ip_info(self, db):
+        """Gets info for the IP found.
+
+           Before attempting API call, it will search the database for an IP match in the previous week.
+        """
+        #data_db = db.query_last_stored_ip_data(self.attempt_ip)
+        #todo https://realpython.com/caching-external-api-requests/
+       # if data_db:
+
+
         return extract_ip_info(self.attempt_ip)