1 files changed, 30 insertions, 31 deletions

diff --git a/cli.py b/cli.py
index f2c8771..bd49707 100644
--- a/cli.py
+++ b/cli.py
@@ -1,53 +1,52 @@
 import argparse
-import json
-from obfuscator.csv_reader import CSVReader
+from obfuscator.read import DataReader
+from obfuscator.write import DataWriter
 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
-    parser = argparse.ArgumentParser(description="gdpr-obfuscator")
-    # Require user to either choose a local file or an S3 object
-    # The user can only choose one of these options or the program will exit
-    # If not provided, the program will exit
+    parser = argparse.ArgumentParser(
+        prog="GDPR-Obfuscator",
+        description="Obfuscate sensitive data stored locally or in an AWS environment",
+    )
+
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose logging"
+    )
+
     loc = parser.add_mutually_exclusive_group(required=True)
-    loc.add_argument("--local")
-    loc.add_argument("--s3")
+    loc.add_argument("-l", "--local", help="Local path to file")
+    loc.add_argument("-s", "--s3", help="URI path to file stored in S3")
 
-    # Require user to provide a list of PII fields to obfuscate
-    # e.g. --pii name email_address
-    # If not provided, the program will exit
-    parser.add_argument("--pii", nargs="+", required=True)
+    parser.add_argument(
+        "-p",
+        "--pii",
+        nargs="+",
+        required=True,
+        help="List of PII fields to obfuscate, separated by spaces",
+    )
 
-    # Parse the arguments
     args = parser.parse_args()
 
-    # Create the CSVReader object
-    reader = CSVReader()
+    log_level = "DEBUG" if args.verbose else "INFO"
+    logger = get_logger("CLI", log_level)
+
+    reader = DataReader(log_level)
 
-    # Read the CSV data based on the user's choice of local or S3
     if args.local and not args.s3:
-        logger.debug("User chose to read CSV from local path")
+        logger.debug("Read data from local path")
         data = reader.read_local(args.local)
-        # For debug purposes, log the data read from the CSV
-        logger.debug("Contents: " + str(data))
     else:
-        logger.debug("User chose to read CSV from S3")
+        logger.debug("Read data from S3")
         data = reader.read_s3(args.s3)
-        logger.debug("Contents: " + str(data))
 
-    # Obfuscate the data based on the user's choice of PII fields
     obfuscated_data = obfuscate(data, args.pii)
-    # For debug purposes, log the obfuscated data as JSON for readability
-    logger.debug("Obfuscated data (JSON): " + json.dumps(obfuscated_data, indent=4))
-    return create_byte_stream(obfuscated_data)
+
+    writer = DataWriter()
+
+    return writer.create_byte_stream(obfuscated_data)
 
 
-# If the script is run directly (as it should be), call the main function
 if __name__ == "__main__":
     main()