Dynamic-Python-Virtual-Environment.py

Dynamic-Python-Virtual-Environment.py

#!/usr/bin/env python

import argparse
import glob
import json
import logging
import os
import pathlib
import shlex
import shutil
import subprocess
import sys
import typing
import abc

logger = logging.getLogger(__name__ if __name__ != "__main__" else "script")

"""
Deletes all of an AWS S3 bucket's versioned files. Optionally perform a dry-run or optionally delete the bucket. 
"""
def main(namespace: argparse.Namespace):
    import boto3
    import botocore.client
    import botocore.config
    import botocore.credentials
    import botocore.exceptions

    verbose: bool = namespace.verbose
    if not verbose:
        logging.getLogger("botocore").setLevel(logging.ERROR)
        logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)

    dry_run: bool = namespace.dry_run

    bucket: str = namespace.bucket_name
    delete_bucket: bool = namespace.include_bucket

    # Configuration for authenticating against aws
    cfg = botocore.config.Config(region_name="us-east-2", signature_version="v4")

    # Establish sts client -- verifies credentials profile
    sts: typing.Union[boto3.BaseClient, typing.Any, None] = None

    try:
        sts = boto3.client("sts", config=cfg)
        logger.debug("Verified AWS Profile")
    except botocore.exceptions.ProfileNotFound as e:
        logger.error("Invalid AWS Profile: %s" % e)
        exit(1)

    # Get aws runtime caller information -- checks if credentials are valid

    try:
        caller = sts.get_caller_identity()
        logger.debug("AWS Caller Identity - Account: %s - User: %s" % (caller.get("Account"), caller.get("Arn")))
    except botocore.exceptions.ProfileNotFound as e:
        logger.error("Invalid AWS Profile. Please Correctly Update \"AWS_PROFILE\" and Try Again: %s" % e)
        exit(1)
    except botocore.exceptions.TokenRetrievalError as e:

        logger.error("Invalid AWS Credentials: %s" % e)
        exit(1)

    # Check if the bucket exists
    s3 = boto3.client("s3", config=cfg)

    try:
        s3.head_bucket(Bucket=bucket)

        logger.info("Found Existing Bucket: %s" % bucket)
    except botocore.exceptions.ClientError as e:
        logger.warning("Unable to Find Bucket: %s" % bucket)
        exit(0)

    resource = boto3.resource("s3", config=cfg)

    target = resource.Bucket(bucket)

    logger.debug("Gathering All Object Version(s). Bucket: %s" % bucket)

    versions: typing.Any

    try:
        versions = target.object_versions.all()
    except Exception as e:
        logger.error("Uncaught Exception: %s" % e, exc_info=True)
        exit(1)

    if dry_run:
        with open("object-versions.json", "w") as outfile:
            json.dump(versions, outfile)

        logger.debug("Wrote Versioned Objects object-versions.json")

        sys.stdout.write("Dry-Run Output Written To: %s\n" % "object-versions.json")

        exit(0)

    logger.debug("Attempting to Delete Object Versions. Bucket: %s" % bucket)

    try:
        target.object_versions.all().delete()
    except Exception as e:
        logger.error("Uncaught Exception: %s" % e, exc_info=True)
        exit(1)

    logger.info("Successfully Deleted All Versioned Object(s)")

    if delete_bucket:
        logger.debug("Attempting to Delete Bucket. Bucket: %s" % bucket)

        try:
            target.delete()
        except Exception as e:
            logger.error("Uncaught Exception: %s" % e, exc_info=True)
            exit(1)

        sys.stdout.write("Deleted Bucket: %s\n" % bucket)
        exit(0)
       
    sys.stdout.write("Deleted All Versioned Object(s). Bucket: %s\n" % bucket)

# e.g. aws-s3-delete-versioned-objects.py
if __name__ == "__main__":
    examples = [
        "  # General usage example that will delete all of an aws s3 bucket's versioned objects",
        "  %(prog)s --bucket-name \"example-bucket\"",
        "",
        "  # Enable debug logging",
        "  %(prog)s --bucket-name \"example-bucket\" --log-level \"DEBUG\"",
        "",
        "  # Write the versioned objects to a file without deleting them",
        "  %(prog)s --bucket-name \"example-bucket\" --dry-run",
        "",
        "  # Additionally delete the bucket after removing all of its files",
        "  %(prog)s --bucket-name \"example-bucket\" --include-bucket",
    ]

    parser = argparse.ArgumentParser(usage="%(prog)s --help", description="", epilog="examples: \n{}".format("\n".join(examples)), formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False)

    # General configuration options
    parser.add_argument("--log-level", type=str, choices=["DEBUG", "INFO", "WARNING", "ERROR"], default="INFO", metavar="level", help="logging verbosity")
    parser.add_argument("--verbose", action="store_true", default=False, help="enable verbose output")

    # Script-specific flags
    parser.add_argument("--bucket-name", type=str, required=True, metavar="value", help="name of the aws s3 bucket")
    parser.add_argument("--include-bucket", action="store_true", default=False, help="additionally delete the aws s3 bucket")
    parser.add_argument("--dry-run", action="store_true", default=False, help="dump the versioned objects without deleting them")

    namespace: argparse.Namespace = parser.parse_args()

    # Setup logging
    logging.basicConfig(format="%(levelname)s - %(name)s - %(filename)s:%(lineno)d - %(message)s", level=
    {"DEBUG": logging.DEBUG, "INFO": logging.INFO, "WARNING": logging.WARNING, "ERROR": logging.ERROR}[namespace.log_level])

    # Dynamically allocate a virtual environment
    if os.getenv("VIRTUAL_ENV") is None or os.getenv("VIRTUAL_ENV").strip() == "":
        logger.warning("No Virtual Environment Found - Dynamically Allocating & Installing Requirement(s)")

        partials = shlex.split("python3 -m venv .venv")

        cmd = partials[0]
        executable = shutil.which(cmd)

        logger.info("Executable: %s" % cmd)
        logger.info("Executable Full-Path: %s" % str(executable))
        logger.info("Command: %s" % " ".join(partials))

        process = subprocess.Popen(partials, executable=executable, shell=False, cwd=os.curdir, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.DEVNULL)

        returncode = process.wait(timeout=30)

        if returncode != 0:
            logger.error("Error While Attempting to Establish Virtual Environment")
            logger.error(process.stderr.read().decode())

            exit(1)

        os.putenv("PATH", pathlib.Path(os.getcwd()).joinpath(".venv", "bin"))

        report = pathlib.Path(os.getcwd()).joinpath("virtual-environment.installation-report.json").resolve()
        os.system(shlex.join(shlex.split("python -m pip --verbose --require-virtualenv --no-cache-dir --no-input --no-color install --report \"%s\" --ignore-installed --force-reinstall boto3" % str(report))))

        py = glob.glob("%s/*" % str(pathlib.Path(os.getcwd()).joinpath(".venv", "lib")), include_hidden=True, recursive=False)

        for path in py:
            packages = pathlib.Path(path).joinpath("site-packages").resolve()

            sys.path.append(str(packages))

    # The primary entrypoint
    main(namespace)