#!/usr/bin/python

import subprocess
import time
import os
import boto
import argparse

from boto.s3.connection import S3Connection
from boto.s3.key import Key
from time import strftime
from boto.s3.lifecycle import Lifecycle, Transition, Rule
from crontab import CronTab


BACKUP_PATH='/tmp/dumper-backups/'

'''

./dumper -a BACKUP -d sparks1  -b 'atp-backups' --aws_key_id=`echo $ATP_AWS_ACCESS_KEY_ID`

./dumper -a RESTORE -d sparks1 -f 'atp_test2015-12-04-13-12-18.sql'  -b 'atp-backups' --aws_key_id=`echo $ATP_AWS_ACCESS_KEY_ID`

'''

def prepare_dir(dir_name):
    if not os.path.exists(dir_name):
        print "[+] Creating directory: %s" % dir_name
        os.makedirs(dir_name)

def establish_connection(access_key, secret_key, bucket_name):
    if access_key is None:
        access_key = os.environ['ATP_AWS_ACCESS_KEY_ID'].strip()
    if secret_key is None:
        secret_key = os.environ['ATP_AWS_SECRET_ACCESS_KEY'].strip()

    conn = S3Connection(access_key, secret_key)
    bucket = conn.get_bucket(bucket_name)

    return bucket


def main():
    parser = argparse.ArgumentParser(prog='dumper',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)


    parser.add_argument('-a', '--action', nargs='?', type=str, default='BACKUP', help='Allowed values: BACKUP or RESTORE')
    parser.add_argument('-s', '--host', nargs='?', type=str, default='localhost', help='host')
    parser.add_argument('-u', '--user', nargs='?', type=str, default='postgres', help='username')
    parser.add_argument('-p', '--password', nargs='?', type=str, default='', help='password')
    parser.add_argument('-d', '--database', nargs='?', type=str, help='database')
    parser.add_argument('-f', '--filename', nargs='?', type=str, help='filename')
    parser.add_argument('-t', '--path', nargs='?', type=str, help='Destination path for backup file')
    parser.add_argument('-e', '--expiry', nargs='?', type=int, help='Expiry in days')
    parser.add_argument('-c', '--schedule', nargs='?', type=int, help='Cron Schedule in hours')
    parser.add_argument('-b', '--bucket', nargs='?', type=str, help='AWS Bucket Name', required=True)

    parser.add_argument('--aws_key_id', nargs='?', type=str, help='AWS Access Key ID')
    parser.add_argument('--aws_secret_key', nargs='?', type=str, help='AWS Secret Access Key')

    args = parser.parse_args()
    # print args

    # set_env_vars()
    bucket = establish_connection(args.aws_key_id, args.aws_secret_key, args.bucket)


    if args.action == "FILES":
        show_keys(bucket)
    elif args.action == "LIFECYCLE":
        if args.expiry:
            set_lifecycle(bucket, args.expiry)
        else:
            show_lifecycle(bucket)
    elif args.action == 'CRON' and args.database and args.bucket:
        if args.schedule:
            # Set schedule
            set_cronjob(args.bucket, args.database, args.schedule)
        else:
            show_cronjob(args.bucket, args.database)

    elif args.action == "BACKUP" and args.database:
        print "[+] Backing up " + args.database
        dump_and_upload_to_s3(args.host, bucket, args.database, args.user)


    elif args.action == "RESTORE" and args.database and args.filename:
        print "RESTORE"
        download_and_restore_from_s3(args.host, bucket, args.filename, 
                                     args.database, args.user, args.path)

    else:
        parser.print_help()

def dump_and_upload_to_s3(host, bucket, database, user):
    location = BACKUP_PATH + "upload/"
    prepare_dir(location)

    filename = database + "-" + time.strftime("%Y-%m-%d-%H-%M-%S") + ".sql"

    dump_database(host, database, user, location + filename)
    upload_to_s3(bucket, location + filename, filename)

def download_and_restore_from_s3(host, bucket, filename, database, user, destination):
    prepare_dir(BACKUP_PATH)

    if destination is None:
        destination = BACKUP_PATH + filename

    download_from_s3(bucket, filename, destination)
    restore_database(host, filename, database, user, destination)


def dump_database(host, database, user, path):
    ps = subprocess.Popen(
          ['pg_dump', '-U', user, '-h', host, '-Fc', database, '-f', path],
          stdout=subprocess.PIPE
      )
    output = ps.communicate()[0]

    print "[+] Database dump location: " + path

def upload_to_s3(bucket, path, filename):
    k = Key(bucket)
    k.key = filename
    k.set_contents_from_filename(path)

    print "[+] File " + filename + " Uploaded"

def download_from_s3(bucket, filename, destination):
    key = bucket.get_key(filename)
    key.get_contents_to_filename(destination)

def restore_database(host, filename, database, user, destination):

    ps = subprocess.Popen(
          ["createdb", database],
          stdout=subprocess.PIPE
      )
    output = ps.communicate()[0]

    ps = subprocess.Popen(["pg_restore", "--verbose", "--clean", "--no-acl", 
                           "--no-owner",  "-h", host, "-U", user, "-d",  database, destination], 
                           stdout=subprocess.PIPE)

    output = ps.communicate()[0]

    print "[+] Finished restoring " + database

def show_keys(bucket):
    for key in bucket.list():
        print key.name.encode('utf-8')

def show_lifecycle(bucket):
    try:
        current = bucket.get_lifecycle_config()
        print current[0].transition
    except:
        print "[+] Lifecycle not yet defined for " + bucket.name + "."

def set_cronjob(bucket, database, hour):
    cron = CronTab(user=True)

    for job in cron.find_comment("-".join(['dumper', bucket, database])):
        print " ".join(["[+] Deleting old cron job for ", database, "at bucket", bucket])
        cron.remove(job)
        cron.write

    cwd = os.getcwd()
    source = ". " + cwd + "/env-vars; "
    command = cwd + '/dumper -a BACKUP -d ' + database + ' -b ' + bucket + " >> dumper-cron.log"
    job  = cron.new(command= source + command, comment="-".join(['dumper', bucket, database]))
    job.setall('0 */' +  str(hour) + ' * * *')
    cron.write()

    print "[+] Added " + job.render()


def show_cronjob(bucket, database):
    cron = CronTab(user=True)
    print " ".join(["[+] Cron jobs for database", database, "at", bucket])
    for job in cron.find_comment("-".join(['dumper', bucket, database])):
        print job

def set_lifecycle(bucket, days):
    to_glacier = Transition(days=days, storage_class='GLACIER')
    rule = Rule('ruleid', '/', 'Enabled', transition=to_glacier)
    lifecycle = Lifecycle()
    lifecycle.append(rule)    

    bucket.configure_lifecycle(lifecycle)

if __name__ == '__main__':
    main()