#!/usr/bin/python3

import socket
import os
import subprocess
import sys
import time
import smtplib
import math
import datetime

dirnum = str(int(math.floor(int(datetime.datetime.now().strftime("%s")) / 86400) % 3))

backupfilesdir = "/backup/" + dirnum
dbname = "intervals" if socket.gethostname().find("b1fr") < 0 else "intervalsdev"
postgresdir = "/var/lib/pgsql/11"
backupdatadir = postgresdir + "/backupdata"
backupwalarchivedir = postgresdir + "/backupwalarchive"
consistentcheck = "consistent recovery state reached"
readycheck = "database system is ready to accept connections"
today = datetime.date.today()
yesterday = datetime.datetime.utcnow() - datetime.timedelta(1)

#db files
dbfiles = [
        "base.tar.gz",
        "pg_wal.tar.gz",
        ]

try:
    #decrypt db backup files
    for dbfile in dbfiles:
        cmd = "gpg --batch --passphrase-fd 0 --output " + postgresdir + "/" + dbfile + "  --decrypt " + backupfilesdir + "/" + dbfile + ".gpg < " + postgresdir + "/.backup_passphrase"
        subprocess.run(cmd, shell=True)

    #uncompress data files into postgres backup data directory (do not use default data dir!)
    subprocess.run(["mkdir", backupdatadir])
    subprocess.run(["chmod", "0700", backupdatadir])
    subprocess.run(["mkdir", backupwalarchivedir])
    subprocess.run(["chmod", "0700", backupwalarchivedir])
    subprocess.run(["tar", "-xzvf", postgresdir + "/base.tar.gz", "-C", backupdatadir])
    subprocess.run(["tar", "-xzvf", postgresdir + "/pg_wal.tar.gz", "-C", backupwalarchivedir])

    #modify postgresql.conf
    subprocess.run(["sed", "-i", "s/shared_buffers/#shared_buffers/g", backupdatadir + "/postgresql.conf"])
    subprocess.run(["sed", "-i", "s/effective_cache_size/#effective_cache_size/g", backupdatadir + "/postgresql.conf"])
    subprocess.run(["sed", "-i", "s/archive_mode/#archive_mode/g", backupdatadir + "/postgresql.conf"])
    subprocess.run(["sed", "-i", "s/maintenance_work_mem/#maintenance_work_mem/g", backupdatadir + "/postgresql.conf"])
    subprocess.run(["sed", "-i", "s/hot_standby =/#hot_standby =/g", backupdatadir + "/postgresql.conf"])

    #modify recovery.conf to take it out of standby mode
    subprocess.run(["sed", "-i", "s/standby_mode/#standby_mode/g", backupdatadir + "/recovery.conf"])
    subprocess.run(["sed", "-i", "s/primary_conninfo/#primary_conninfo/g", backupdatadir + "/recovery.conf"])
    #change where recovery reads its wal files from
    subprocess.run(["sed", "-i", "s/restore_command/#restore_command/g", backupdatadir + "/recovery.conf"])
    subprocess.run(["sed", "-i", "s/archive_cleanup_command/#archive_cleanup_command/g", backupdatadir + "/recovery.conf"])
    with open(backupdatadir + "/recovery.conf", "w+") as recoveryfile:
        recoveryfile.write("restore_command = 'cp -i " + backupwalarchivedir + "/%f %p'\n")
        recoveryfile.write("archive_cleanup_command = '/usr/pgsql-11/bin/pg_archivecleanup " + backupwalarchivedir + " %r'\n")

    #start postgresql
    subprocess.run(["/usr/pgsql-11/bin/pg_ctl", "start", "-D", backupdatadir, "-t", "150"])

    #check that recovery has completed

    i = 0
    while i < 300:
        rout = subprocess.run(["grep", consistentcheck, backupdatadir + "/log/postgresql-" + str(today) + ".log"],stdout=subprocess.PIPE)
        print('Waiting for recovery to complete...')
        if rout.stdout.decode('utf-8').find(consistentcheck) > 0:
            print(consistentcheck)
            break
        time.sleep(1)
        i += 1

    #uncomment the line below to exit this script and leave the db up and running
    #sys.exit(2)

    #run a query against the database to make sure we have recent data
    rout =subprocess.run(["psql " + dbname + " -t -c 'SELECT create_time FROM mail_queue ORDER BY id DESC LIMIT 1'"],shell=True,stdout=subprocess.PIPE)

    #read the date from query
    backupdate = rout.stdout.decode('utf-8').strip()

    #pause here for a little bit to let logging catch up
    time.sleep(10)

    #read the status file
    with open(backupdatadir + "/log/postgresql-" + str(today) + ".log", 'r') as statusfileio:
        backupstatus = statusfileio.read()

    #stop the database, destroy the files we created, and stop the container so we have a clean slate next time
    subprocess.run(["/usr/pgsql-11/bin/pg_ctl", "stop", "-D", backupdatadir])
    subprocess.run(["rm", "-rf", backupdatadir])
    subprocess.run(["rm", "-rf", backupwalarchivedir])

    #delete unencrypted files
    for dbfile in dbfiles:
        os.remove(postgresdir + "/" + dbfile)

    #convert date string to datetime object
    backupdate = datetime.datetime.strptime(backupdate, '%Y-%m-%d %H:%M:%S')
        
    #if yesterday's date is not present in the date file, alert the authorities!
    if backupdate < yesterday:
        msg = "From: backup@pelagodesign.com\r\n"
        msg = msg + "To: backup@pelagodesign.com\r\n"
        #msg = msg + "Subject: PostgreSQL Backup Test Failed\r\n"
        msg = msg + "Subject: Staging: PostgreSQL Results (fail)\r\n"
        msg = msg + "\r\n"
        msg = msg + "The backup did not contain data for yesterday. Please investigate.\n"
        msg = msg + "\nDate found: "+backupdate.strftime('%Y-%m-%d %H:%M:%S')+" UTC\n"
        msg = msg + "Logs:\n"+backupstatus
        print(msg)

        #server = smtplib.SMTP("192.168.132.78")
        server = smtplib.SMTP("localhost")
        server.sendmail("backup@pelagodesign.com", "backup@pelagodesign.com", msg.encode("utf-8"))
        server.quit()
    else:
        msg = "From: backup@pelagodesign.com\r\n"
        msg = msg + "To: backup@pelagodesign.com\r\n"
        #msg = msg + "Subject: PostgreSQL Backup Test Success!\r\n"
        msg = msg + "Subject: Staging: PostgreSQL Results (success)\r\n"
        msg = msg + "\r\n"
        msg = msg + "\nDate found: "+backupdate.strftime('%Y-%m-%d %H:%M:%S')+" UTC\n"
        msg = msg + "Logs:\n"+backupstatus
        print(msg)

        #server = smtplib.SMTP("192.168.132.78")
        server = smtplib.SMTP("localhost")
        server.sendmail("backup@pelagodesign.com", "backup@pelagodesign.com", msg.encode("utf-8"))
        server.quit()

except SystemExit:
    sys.exit(2)

except:
    errormessage = str(sys.exc_info()[0]) + "\n" + str(sys.exc_info()[1]) + "\n" + str(sys.exc_info()[2])

    msg = "From: backup@pelagodesign.com\r\n"
    msg = msg + "To: backup@pelagodesign.com\r\n"
    #msg = msg + "Subject: PostgreSQL Backup Test Failed\r\n"
    msg = msg + "Subject: Staging: PostgreSQL Results (error)\r\n"
    msg = msg + "\r\n"
    msg = msg + "The backup cron had errors. Please investigate.\n"
    msg = msg + "\Error: "+errormessage+"\n"
    print(msg)

    #server = smtplib.SMTP("192.168.132.78")
    server = smtplib.SMTP("localhost")
    server.sendmail("backup@pelagodesign.com", "backup@pelagodesign.com", msg.encode("utf-8"))
    server.quit()
