EStateGetFromS3.py

#!/usr/bin/env python3

"""Retrieve eStatement from AWS S3

Only the correct content is streamed to stdout. Any error must be
streamed down to stderr descriptor.

On failure, this script retries aws s3 api call for the number provided with
--s3-retries option argument. Default is MAX_RETRIES_DEFAULT.

Script Output:
    Status code:
        0, if no error
        1, if any error, exceptions

    Content (see usage for more detail):
        PDF binary content, OR
        S3 Presigned URL - has a short expiration duration
                       (ENV: AWS_S3_PRESIGNURL_EXPIRATION_SECONDS;
                       default: 120 seconds)

References:
[1] https://docs.aws.amazon.com/cli/latest/reference/s3/presign.html
[2] https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.generate_presigned_url

"""


import sys
import boto3
import os
import logging

from argparse import ArgumentParser

format=("[%(levelname)-5s] %(asctime)-15s ({}|%(funcName)s|%(lineno)d) "
        ":: %(message)s").format(__file__)
LOGGER = logging.getLogger(__name__)

LOGGER.setLevel(logging.ERROR)
logging_handler = logging.StreamHandler()
logging_handler.setFormatter(logging.Formatter(format))
LOGGER.addHandler(logging_handler)


# [Prod]
AWS_REGION = 'us-east-1'
S3_STMNT_BUCKET = "homecu.stmnts"

# [Dev] Use the following settings; simply uncomment the following and edit
# as necessary during development. Work with infra team to copy necessary
# prefixes and objects into this bucket and set appropriate permissions.
# AWS_REGION = 'us-east-2'
# S3_STMNT_BUCKET = "homecu.dev.stmnts"

# Maximum retries on aws api call failure
MAX_RETRIES_DEFAULT = 2
# expiration period for pre-signed urls (in seconds)
PRESIGN_URL_EXPIRATION_SECONDS = int(os.environ.get(
    'AWS_S3_PRESIGNURL_EXPIRATION_SECONDS', 120))


def get_parser():
    """Prepare argument parser.

    Returns:
        parser -- ArgumentParser object
    """
    parser = ArgumentParser(
        description=("Get PDF EStatement Content from S3 (buffered"
        "binary content OR s3 presigned url).")
    )

    parser.add_argument("-f", "--file-estmnt",
                        help="Full path to estatement file for this period",
                        required=True
                        )

    parser.add_argument("-p", "--presignedurl", action = "store_true",
                        help=("stream s3 presign url to stdout; "
                              "otherwise streams actual pdf content")
                        )
    parser.add_argument("-r", "--s3-retries",
                        help="Maximum number of retries on failure",
                        type=int,
                        default=MAX_RETRIES_DEFAULT)
    return parser


def cleanup_temp_file(tmp_file):
    """Cleanup temporarily downloaded estatement file"""
    if os.path.exists(tmp_file):
        os.remove(tmp_file)


def return_presigned_url(s3_client, object_key):
    """Return an object's pre-signed"""
    url = ""
    try:
        url = s3_client.generate_presigned_url(
            ClientMethod='get_object',
            Params={
                "Bucket" : S3_STMNT_BUCKET,
                "Key" : object_key
            },
            ExpiresIn=PRESIGN_URL_EXPIRATION_SECONDS
    )
    except Exception as e:
        LOGGER.error("Estatement `%s` download failed." % object_key)
        raise
    else:
        # return s3 presigned url
        print(url)


def return_pdf_content(s3_client, s3_file_name):
    """Download pdf from S3 and return stream content to stdout"""
    tmp_file = os.path.join("/tmp", s3_file_name.replace("/", "_"))
    try:
        # download pdf object from S3
        s3_client.download_file(S3_STMNT_BUCKET, s3_file_name, tmp_file)
    except Exception as e:
        cleanup_temp_file(tmp_file)
        LOGGER.error("Estatement `%s` download failed." % s3_file_name)
        raise
    else:
        # return pdf content
        with open(tmp_file, 'rb') as f_read:
           sys.stdout.buffer.write(f_read.read())
        cleanup_temp_file(tmp_file)


def main(argv):
    """Retrieve eStatement from AWS S3 estatement storing bucket"""

    parser=get_parser()
    args = parser.parse_args(argv)
    success_status = False

    file_parts = args.file_estmnt.split("/")[1:]

    if len(file_parts) != 7:
        LOGGER.error(("Requested path to estatement '%s' is invalid; "
                     "Eg. correct path: "
                     "'/home/CUCODE/stmnt/pdf/yyyymmM/1/1XXXX.pdf" %
                     args.file_estmnt))
        return success_status

    cucode = file_parts[1].strip()
    s3_file_name = os.path.join(file_parts[1],
                                file_parts[2],
                                file_parts[3],
                                file_parts[4],
                                file_parts[5],
                                file_parts[6])
    max_retries = args.s3_retries
    tries = 0

    while tries < max_retries and not success_status:
        tries += 1
        try:
            client = boto3.client('s3', region_name=AWS_REGION)

            if args.presignedurl:
                return_presigned_url(client, s3_file_name)
            else:
                return_pdf_content(client, s3_file_name)

        except Exception as e:
            LOGGER.error(("[%s] Estatement Retrieval from AWS S3 failed "
                          "(try# %s/%s): %s") %  (cucode.upper(), tries,
                                                  max_retries, e))
        else:
            success_status = True

    return success_status


def run():
    """Script entrypoint"""
    sys.exit(main(sys.argv[1:]))


if __name__ == "__main__":
    run()