diff options
| -rw-r--r-- | cli.py | 12 | ||||
| -rw-r--r-- | obfuscator/csv_reader.py | 38 | ||||
| -rw-r--r-- | obfuscator/logger.py | 41 | ||||
| -rw-r--r-- | obfuscator/utils.py | 27 |
4 files changed, 67 insertions, 51 deletions
@@ -5,9 +5,6 @@ from obfuscator.obfuscate import obfuscate from obfuscator.logger import get_logger from obfuscator.csv_writer import create_byte_stream -# Create the logger -logger = get_logger("CLI") - def main(): # Create an argument parser @@ -42,11 +39,13 @@ def main(): args = parser.parse_args() # If the user chose verbose logging, set the logger to debug - if args.verbose: - logger.setLevel("DEBUG") + log_level = "DEBUG" if args.verbose else "INFO" + + # Create the logger + logger = get_logger("CLI", log_level) # Create the CSVReader object - reader = CSVReader() + reader = CSVReader(log_level) # Read the CSV data based on the user's choice of local or S3 if args.local and not args.s3: @@ -56,6 +55,7 @@ def main(): logger.debug("Contents: " + str(data)) else: logger.debug("User chose to read CSV from S3") + data = reader.read_s3(args.s3) logger.debug("Contents: " + str(data)) diff --git a/obfuscator/csv_reader.py b/obfuscator/csv_reader.py index 8fdf26f..f8dd7d3 100644 --- a/obfuscator/csv_reader.py +++ b/obfuscator/csv_reader.py @@ -4,10 +4,7 @@ import boto3 import os from typing import List, Dict from obfuscator.logger import get_logger -from obfuscator.utils import get_s3_path - -# Create the logger -logger = get_logger("CSVREADER") +from obfuscator.utils import Utilities # Putting the CSV reading components into a class may seem like overkill # for a simple script, but it allows for better organization and scalability. @@ -22,14 +19,18 @@ class CSVReader: the project completion, support for JSON/Parquet files will be added. """ - @staticmethod - def read_local(path) -> List[Dict[str, str]]: + def __init__(self, log_level=None): + self.log_level = log_level + # Create the logger + self.logger = get_logger("CSVREADER", log_level) + + def read_local(self, path) -> List[Dict[str, str]]: """ A method to read a local CSV file and return the data as a list of dictionaries. """ # Log the path of the file being read for debugging - logger.debug(f"Reading local CSV from: {path}") + self.logger.debug(f"Reading local CSV from: {path}") # Attempt to read the file and return the data as a list of dictionaries # However, if the file isn't found or there is a generic exception, log @@ -39,38 +40,38 @@ class CSVReader: reader = csv.DictReader(f) return [dict(row) for row in reader] except FileNotFoundError: - logger.error(f"File not found: {path}") + self.logger.error(f"File not found: {path}") raise except Exception as e: - logger.error(f"Error reading file: {e}") + self.logger.error(f"Error reading file: {e}") - @staticmethod - def read_s3(path) -> List[Dict[str, str]]: + def read_s3(self, path) -> List[Dict[str, str]]: """ A method to read an S3 object containing CSV data and return the data as a list of dictionaries. """ - bucket, key = get_s3_path(path) - logger.debug(f"Reading S3 CSV from: {bucket}/{key}") + utils = Utilities(self.log_level) + bucket, key = utils.get_s3_path(path) + self.logger.debug(f"Reading S3 CSV from: {bucket}/{key}") # If LOCALSTACK=TRUE, use the localstack endpoint for testing if os.getenv("LOCALSTACK", "FALSE").upper() == "TRUE": localstack_endpoint = "http://localhost.localstack.cloud:4566" - logger.debug("Using LocalStack endpoint for S3") + self.logger.debug("Using LocalStack endpoint for S3") client = boto3.client( "s3", endpoint_url=localstack_endpoint, aws_access_key_id="dummy", aws_secret_access_key="dummy", ) - logger.debug(f"endpoint_url: {localstack_endpoint}") + self.logger.debug(f"endpoint_url: {localstack_endpoint}") else: client = boto3.client("s3") try: # Attempt to read the S3 object and return the data as a list of dictionaries response = client.get_object(Bucket=bucket, Key=key) - logger.info("S3 object read successfully") + self.logger.info("S3 object read successfully") # Read and decode the content content = response["Body"].read().decode("utf-8") # Even though the read_string method was only created for testing, @@ -78,11 +79,10 @@ class CSVReader: return CSVReader.read_string(content) # TODO: Add more specific exceptions to catch except Exception as e: - logger.error(f"Error reading S3 object: {e}") + self.logger.error(f"Error reading S3 object: {e}") raise - @staticmethod - def read_string(content: str) -> List[Dict[str, str]]: + def read_string(self, content: str) -> List[Dict[str, str]]: """ A method to read CSV data from a string and return the data as a list of dictionaries. diff --git a/obfuscator/logger.py b/obfuscator/logger.py index 649dad7..2c5b988 100644 --- a/obfuscator/logger.py +++ b/obfuscator/logger.py @@ -1,24 +1,37 @@ import logging import os +from enum import Enum -def get_logger(name: str) -> logging.Logger: - logger = logging.getLogger(name) +class LogLevel(Enum): + DEBUG = logging.DEBUG + INFO = logging.INFO + WARNING = logging.WARNING + ERROR = logging.ERROR + CRITICAL = logging.CRITICAL + - if not logger.hasHandlers(): - if os.getenv("DEBUG", "FALSE").upper() == "TRUE": - log_level = logging.DEBUG - else: - log_level = logging.INFO +def get_logger(name: str, level: LogLevel = LogLevel.INFO) -> logging.Logger: - logger.setLevel(log_level) + if isinstance(level, str): + try: + level = LogLevel[level.upper()] + except KeyError: + raise ValueError( + f"Invalid log level '{level}'. Choose from: {', '.join(l.name for l in LogLevel)}" + ) + + logger = logging.getLogger(name) - handler = logging.StreamHandler() - formatting = logging.Formatter( - "[%(asctime)s] - %(levelname)s::%(name)s - %(message)s" - ) - handler.setFormatter(formatting) + if logger.hasHandlers(): + logger.handlers.clear() - logger.addHandler(handler) + handler = logging.StreamHandler() + logger.setLevel(level.value) + formatting = logging.Formatter( + "[%(asctime)s] - %(levelname)s::%(name)s - %(message)s" + ) + handler.setFormatter(formatting) + logger.addHandler(handler) return logger diff --git a/obfuscator/utils.py b/obfuscator/utils.py index 1d1c3fe..81eb04a 100644 --- a/obfuscator/utils.py +++ b/obfuscator/utils.py @@ -1,15 +1,18 @@ # Utility functions from obfuscator.logger import get_logger -# Create the logger -logger = get_logger("UTILS") - - -def get_s3_path(uri): - parts = uri.replace("s3://", "").split("/") - logger.debug(f"Parts: {parts}") - bucket = parts.pop(0) - logger.debug(f"Bucket: {bucket}") - key = "/".join(parts) - logger.debug(f"Key: {key}") - return bucket, key + +class Utilities: + + def __init__(self, logger=None): + # Create the logger + self.logger = get_logger("UTILITIES", logger) + + def get_s3_path(self, uri): + parts = uri.replace("s3://", "").split("/") + self.logger.debug(f"Parts: {parts}") + bucket = parts.pop(0) + self.logger.debug(f"Bucket: {bucket}") + key = "/".join(parts) + self.logger.debug(f"Key: {key}") + return bucket, key |