4
\$\begingroup\$

I have server side log files written to various directories throughout the course of a week. I have the need to optionally compress, and archive these files to AWS.

The script I have come up with takes a config file defining the include glob to look for matches, and offers optional removal of the source files once the process has completed.

python aws_bkup.py path/to/config.cfg

A sample config file can be found in the project directory: https://github.com/twanas/aws-bkup

I would be most appreciative of feedback on the code style and areas it can be improved.

from __future__ import print_function
import re
import errno
import os
from glob import glob
from os.path import join, basename, splitext
from os import environ, remove
from shutil import copy, rmtree
from uuid import uuid4
import gzip
from configparser import ConfigParser
from datetime import date, timedelta
import subprocess


def gz(src, dest):
    """ Compresses a file to *.gz
        Parameters
        ----------
        src: filepath of file to be compressesd
        dest: destination filepath
    """

    filename = splitext(basename(src))[0]
    destpath = join(dest, '{}.gz'.format(filename))

    blocksize = 1 << 16     #64kB

    with open(src) as f_in:
        f_out = gzip.open(destpath, 'wb')
        while True:
            block = f_in.read(blocksize)
            if block == '':
                break
            f_out.write(block)
        f_out.close()


def aws_sync(src, dest):
    """ Synchronise a local directory to aws
    Parameters
    ----------
    src: local path
    dest: aws bucket
    """
    cmd = 'aws s3 sync {} {}'.format(src, dest)
    push = subprocess.call(cmd, shell=True)


def today():
    """ Returns a string format of today's date """
    return date.today().strftime('%Y%m%d')


def fwe():
    """ Returns a string format of the next friday's date """
    d = date.today()
    while d.weekday() != 4:
        d += timedelta(1)
    return d


def regex_match(string, pattern):
    """ Returns if there is a match between parameter and regex pattern """
    pattern = re.compile(pattern)
    return pattern.match(string)


def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def aws_bkup(section, include, exclude, s3root, categorize_weekly=True, compress=True, remove_source=True):
    """ Transfers a backup of any local files matching the user's criteria to AWS.
    Parameters
    ----------
    include: regex pattern to use for the file inclusion(s)
    exclude: regex pattern to use for the file exclusion(s)
    s3root: AWS root in which to send the backup
    categorize_weekly: switch between daily and weekly folder groupings
    compress: switch to compress outbound files to AWS
    """

    folder = '{}'.format(fwe() if categorize_weekly else today())
    tmp_root = join('/tmp', str(uuid4()))
    tmp_dir = join(tmp_root, folder)

    mkdir_p(tmp_dir)

    for file in glob(include):

        if regex_match(file, exclude):
            continue

        print('Processing: {}'.format(file))

        if compress:
            gz(file, tmp_dir)
        else:
            copy(file, tmp_dir)

        if remove_source:
            remove(file)

    aws_dest = join(s3root, section)

    print('Syncronizing {} to s3'.format(tmp_dir))
    aws_sync(tmp_root, aws_dest)

    if os.path.exists(tmp_root):
        rmtree(tmp_root)

    print('Done')


if __name__ == "__main__":

    import sys
    args = sys.argv

    if len(args) < 2:
        print("Usage: python -m aws-bkup /path/to/config.cfg")
        sys.exit()

    config = ConfigParser()
    config.read(args[1])

    environ['AWS_ACCESS_KEY_ID'] = config.get('aws', 'access_id')
    environ['AWS_SECRET_ACCESS_KEY'] = config.get('aws', 'secret_key')
    environ['AWS_DEFAULT_REGION'] = config.get('aws', 'region')

    for section in config.sections():
        if section != 'aws':
            print('Starting {}'.format(section))
            aws_bkup(
                section,
                config.get(section, 'include'),
                config.get(section, 'exclude'),
                config.get('aws', 's3root'),
                config.getboolean(section, 'categorize_weekly'),
                config.getboolean(section, 'compress'),
                config.getboolean(section, 'remove_source')
            )
\$\endgroup\$
2
\$\begingroup\$

After a quick read-through, I’ve spotted two items:


with not used for f_out

The code:

with open(src) as f_in:
    f_out = gzip.open(destpath, 'wb')
    #...
    f_out.close()

should be replaced with:

with open(arc) as f_in, gzip.open(destpath, 'wb') as f_out:
    #...

Reg-ex pattern repeatedly compiled

The function regex_match() takes a string and compiles it to a pattern, and then matches a string to that pattern. The same pattern string is repeatedly passed to regex_match. This string should be compiled to a pattern by the caller, and the resulting pattern reused for each match. This means the calls to regex_match could be replaced by exclude_pattern.match(file)


Argument quoting

If src or dest contain spaces, this command may become confused.

cmd = 'aws s3 sync {} {}'.format(src, dest)
push = subprocess.call(cmd, shell=True)

Since you are using the shell=True argument, it may also be a vector for arbitrary command injection!

Instead of formatting the command into a string, with proper quoting, and requiring the .call() command to parse it, you can simply pass in an array of arguments to the call. No need to worry about spaces or proper escaping/quoting -- and arbitrary command injection becomes much harder:

cmd = ['aws', 's3', 'sync', src, dest]
push = subprocess.call(cmd, shell=True)

Additional notes:

push is neither returned or used.

Also, while subprocess.call(...) is still acceptable, as of Python 3.5 subprocess.run(...) is the preferred interface.


More later

\$\endgroup\$
  • \$\begingroup\$ Thank you for the feedback @AJNeufeld. I have made the recommended adjustments. I appreciate your help. \$\endgroup\$ – awaugh Aug 22 '18 at 9:02

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.