#!/usr/bin/env python2
# -*- coding: utf-8 -*- 

#   Bacula(R) - The Network Backup Solution

#   Copyright (C) 2000-2023 Kern Sibbald

#   The original author of Bacula is Kern Sibbald, with contributions
#   from many others, a complete list can be found in the file AUTHORS.

#   You may use this file and others of this release according to the
#   license defined in the LICENSE file, which includes the Affero General
#   Public License, v3.0 ("AGPLv3") and some additional permissions and
#   terms pursuant to its AGPLv3 Section 7.

#   This notice must be preserved when any source code is
#   conveyed and/or propagated.

#   Bacula(R) is a registered trademark of Kern Sibbald.

#
#  Routines for aws S3 cloud driver
#
#  Written by Norbert Bizet, August MMXVIII
#
#

import sys, os, json, time
from subprocess import Popen, PIPE, call
from multiprocessing import Pool, cpu_count
from inspect import stack
import logging, traceback


# RETRY DOWNLOAD
RETRY_DOWNLOAD = 0x0D

def vol_ls():
   try:
      logging.info("enter vol_ls")
      cmd = ["aws", "s3", "ls", cloud_path]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True)
      output,err = proc.communicate()
      logging.debug("vol_ls proc communicate output:{0} , err:{1}".format(output,err))
      # sort out only exception errors since progress is reported into stderr
      if err:
         logging.error("vol_ls got error {0}".format(err))
         sys.stderr.write(err)
      if output:
         # expected output format will be "VolName1\nVolName2\nVolName3\n"
         output = "\n".join(list(filter(None, [line.rsplit(' ',2)[2].strip() for line in output.replace(cloud_path, '').replace('/','').splitlines()]))) + "\n"
         # forward out stds
         logging.info("vol_ls got ouput")
         logging.info("vol_ls outputing {0}".format(output))
         sys.stdout.write(output)
      if proc.returncode == 1:
         # possible to ls unexisting path. In this case, return code will be 1.
         return 0
      return proc.returncode
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1

def ls():
   try:
      logging.info("enter ls")
      cmd = ["aws", "s3", "ls", os.path.join(cloud_path, volume, part)]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True)
      output,err = proc.communicate()
      logging.debug("ls proc communicate output:{0} , err:{1}".format(output,err))
      # sort out only exception errors since progress is reported into stderr
      if err:
         logging.error("ls got error {0}".format(err))
         sys.stderr.write(err)
      if output:
         logging.info("ls got ouput")
         # parse and format output
         #mtime
         mtimes = [int(time.mktime(time.strptime(line.rsplit(' ',2)[0].strip(),'%Y-%m-%d %H:%M:%S'))) for line in output.splitlines()]
         #size
         sizes = [line.rsplit(' ',2)[1].strip() for line in output.splitlines()]
         #names
         names = [line.rsplit(' ',2)[2].strip() for line in output.splitlines()]
         # expected output format will be "name:Name1,mtime:Time1,size:Size1\nname:Name2,mtime:Time2,size:Size2\nname:Name3,mtime:Time3,size:Size3\n"
         output = "\n".join(["name:{0},mtime:{1},size:{2}".format(n,t,s) for n, t, s in zip(names, mtimes, sizes)]) + "\n"
         # forward out stds
         logging.info("ls outputing {0}".format(output))
         sys.stdout.write(output)
      return proc.returncode
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1

def download():
   try:
      logging.info("enter download")
      cmd = ["aws", "s3", "cp", os.path.join(cloud_path, volume, part), "-", "--only-show-errors"]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stderr=PIPE, universal_newlines=True)
      output,err = proc.communicate()
      logging.debug("download proc communicate output:{0} , err:{1}".format(output,err))
      if err and err.find("An error occurred (InvalidObjectState) when calling the GetObject operation") != -1:
         restore()
         return RETRY_DOWNLOAD
      return proc.returncode
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1

def delete():
   try:
      logging.info("enter delete")
      cmd = ["aws", "s3", "rm", os.path.join(cloud_path, volume, part), "--only-show-errors"]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True)
      output,err = proc.communicate()
      logging.debug("delete proc communicate output:{0} , err:{1}".format(output,err))
      # sort out only exception errors since progress is reported into stderr (yuck!)
      if err:
         logging.error("delete got error {0}".format(err))
         sys.stderr.write(err)
      if output:
         logging.info("delete got ouput {0}".format(output))
         sys.stdout.write(output)
      return proc.returncode
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1

def upload():
   try:
      logging.info("enter upload")
      cmd = ["aws", "s3", "cp", "-", os.path.join(cloud_path, volume, part),
         "--storage-class", objects_default_tier, "--only-show-errors"]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True)
      output,err = proc.communicate()
      logging.debug("upload proc communicate output:{0} , err:{1}".format(output,err))
      # sort out only exception errors since progress is reported into stderr (yuck!)
      if err:
         sys.stderr.write(err)
      ret = proc.returncode
      cmd = ["aws", "s3", "ls", os.path.join(cloud_path, volume, part)]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True)
      output,err = proc.communicate()
      logging.debug("ls proc communicate output:{0} , err:{1}".format(output,err))
      # sort out only exception errors since progress is reported into stderr
      if err:
         logging.error("ls got error {0}".format(err))
         sys.stderr.write(err)
      if output:
         logging.info("ls got ouput")
         # parse and format output
         #mtime
         mtimes = [int(time.mktime(time.strptime(line.rsplit(' ',2)[0].strip(),'%Y-%m-%d %H:%M:%S'))) for line in output.splitlines() if line.endswith(part)]
         #size
         sizes = [line.rsplit(' ',2)[1].strip() for line in output.splitlines() if line.endswith(part)]
         #names
         names = [line.rsplit(' ',2)[2].strip() for line in output.splitlines() if line.endswith(part)]
         # expected output format will be "name:Name1,mtime:Time1,size:Size1\nname:Name2,mtime:Time2,size:Size2\nname:Name3,mtime:Time3,size:Size3\n"
         output = "\n".join(["name:{0},mtime:{1},size:{2}".format(n,t,s) for n, t, s in zip(names, mtimes, sizes)]) + "\n"
         # forward out stds
         logging.info("ls outputing {0}".format(output))
         sys.stdout.write(output)
      ret = ret + proc.returncode
      return ret
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1

def move():
   try:
      logging.info("enter move from {0} to {1}".format(os.path.join(cloud_path, volume, part), os.path.join(cloud_path, volume, local_part)))
      cmd = ["aws", "s3", "mv",
         os.path.join(cloud_path, volume, part),
         os.path.join(cloud_path, volume, local_part), "--only-show-errors"]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True)
      output,err = proc.communicate()
      logging.debug("move proc communicate output:{0} , err:{1}".format(output,err))
      if not err:
         sys.stdout.write("{0}\0".format(local_part))
         return 0
      else:
         if err.find("Key \"{0}\" does not exist".format(os.path.join(volume, part))) != -1:
            logging.debug("move cant find source {0}. OK.".format(part))
            return 0
         logging.error("move got error {0}".format(err))
         sys.stderr.write(err)
      return proc.returncode
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1

def restore():
   try:
      logging.info("enter restore {0}".format(os.path.join(cloud_path, volume, part)))
      cmd = ["aws", "s3api", "restore-object", 
         "--restore-request", "{{\"Days\": {0}, \"GlacierJobParameters\" : {{\"Tier\" : \"{1}\"}}}}".format(transfer_retention_days,transfer_priority),
         "--bucket", bucket,
         "--key", os.path.join(volume, part)]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True) 
      output,err = proc.communicate()
      logging.debug("restore proc communicate output:{0} , err:{1}".format(output,err))
      if err:
         logging.error("restore got error {0}".format(err))
         sys.stderr.write(err)
      if output:
         logging.info("restore got ouput {0}".format(output))
         sys.stdout.write(output)
      return proc.returncode
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1

def wait_on_restore():
   try:
      logging.info("enter wait_on_restore {0}".format(os.path.join(cloud_path, volume, part)))
      cmd = ["aws", "s3api", "head-object", 
         "--bucket", bucket,
         "--key", os.path.join(volume, part)]
      if endpoint_url:
         cmd += ["--endpoint-url", endpoint_url]
      proc = Popen( cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True) 
      output,err = proc.communicate()
      logging.debug("wait_on_restore proc communicate output:{0} , err:{1}".format(output,err))
      if err:
         logging.error("wait_on_restore got error {0}".format(err))
         sys.stderr.write(err)
         return 1
      if output:
         logging.info("wait_on_restore got ouput {0}".format(output))
         j = json.loads(output)
         if "Restore" in j and j["Restore"]:
            if j["Restore"]=='ongoing-request="true"':
               logging.info("Ongoing restore detected")
               sys.stdout.write("WOR-INPROGRESS")
            elif j["Restore"]=='ongoing-request="false"':
               logging.info("NO Ongoing restore detected")
               sys.stdout.write("WOR-DONE")
      return 0
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
   return 1


if __name__ == '__main__':
   # initialize the return code with a weird value
   ret = 10

   if ('CLOUD_DEBUG' in os.environ) and os.environ['CLOUD_DEBUG']:
      logging.basicConfig(filename='@working_dir@/aws_cloud_driver.log', level=logging.DEBUG, filemode='a', format='%(asctime)-15s-%(process)d- %(levelname)s -%(message)s')
   else:
      logging.basicConfig(level=logging.CRITICAL)

   try:
      if len(sys.argv) < 3:
         sys.stderr.write("google_driver: invalid number of parameters\0.")
         ret = 1

      fct=sys.argv[1]
      volume=sys.argv[2]
      if len(sys.argv) > 3:
         part=sys.argv[3]
      else:
         part=''
      if len(sys.argv) > 4:
         local_part=sys.argv[4]
      else:
         local_part='*None*'
      logging.info("--{1} {2} {3} {4} --".format(sys.argv[0], fct, volume, part, local_part))

      bucket=os.environ['CLOUD_BUCKET']
      cloud_path = "s3://{0}".format(bucket)
      access_key=os.environ['CLOUD_ACCESS_KEY']
      secret_key=os.environ['CLOUD_SECRET_KEY']
      region=os.environ['CLOUD_REGION']
      protocol=os.environ['CLOUD_PROTOCOL']
      uri_type=os.environ['CLOUD_URI_TYPE']
      endpoint_url=os.environ['CLOUD_BLOB_ENDPOINT']

      transfer_priority_switcher = {
         "0":"Expedited",
         "1":"Standard",
         "2":"Bulk"
      }
      t_prio = os.environ['CLOUD_TRANSFER_PRIORITY']
      transfer_priority = transfer_priority_switcher[t_prio]
      transfer_retention_days = os.environ['CLOUD_TRANSFER_RETENTION_DAYS']
      objects_default_tier_switcher = {
         "0":"STANDARD",
         "1":"STANDARD_IA",
         "2":"INTELLIGENT_TIERING",
         "3":"ONEZONE_IA",
         "4":"GLACIER_IR",
         "5":"GLACIER",
         "6":"DEEP_ARCHIVE",
         "7":"REDUCED_REDUNDANCY"
      }
      t_tier_idx = os.environ['CLOUD_OBJECTS_DEFAULT_TIER']
      objects_default_tier = objects_default_tier_switcher[t_tier_idx]

      os.environ['AWS_DEFAULT_REGION'] = region
      os.environ['AWS_ACCESS_KEY_ID'] = access_key
      os.environ['AWS_SECRET_ACCESS_KEY'] = secret_key

      logging.info("bucket {0}, cloud_path {1}, access_key {2}, secret_key {3}, region {4}, protocol {5}, uri_type {6}, transfer_prio {7}, tranfer_ret {8}, default_tier {9}, endpoint_url {10}".format(bucket, cloud_path, access_key, "XXX", region, protocol, uri_type, transfer_priority, transfer_retention_days, objects_default_tier, endpoint_url))

      switcher = {
         "vol_ls":vol_ls,
         "ls":ls,
         "download":download,
         "delete":delete,
         "upload":upload,
         "move":move,
         "wait_on_restore":wait_on_restore
      }

      if sys.argv[1] in switcher:
         ret = switcher[sys.argv[1]]()
         if ret is 0:
            logging.info("{0} returned {1}".format(fct, ret))
         else:
            logging.error("{0} returned {1}".format(fct, ret))
      else:
         sys.stderr.write('unsupported function {0}\n'.format(sys.argv[1]))
         ret = 3
   except Exception as e:
      exc = traceback.format_exception_only(type(e), e)[0]
      sys.stderr.write(exc)
      logging.error(traceback.format_exc())
      ret = 4

   #only once outside of the try-catch statement
   sys.exit(ret)