#!/usr/bin/python3
# AWS EC2 HibInit Agent. This agent does several things:
# 1. Upon startup it checks for sufficient swap space to allow hibernate and fails
#    if it's present but there's not enough of it.
# 2. If there's no swap space, it creates it and launches a background thread to
#    touch all of its blocks to make sure that EBS volumes are pre-warmed.
# 3. It updates the offset of the swap file in the kernel using SNAPSHOT_SET_SWAP_AREA ioctl.
#
# This file is compatible both with Python 2 and Python 3
import argparse
import array
import atexit
import ctypes as ctypes
import fcntl
import mmap
import os, signal
import struct
import sys
import syslog
import math
from subprocess import check_call, check_output, STDOUT
from threading import Thread
from math import ceil
from time import sleep


try:
    from urllib.request import urlopen, Request
except ImportError:
    from urllib2 import urlopen, Request, HTTPError

try:
    from ConfigParser import ConfigParser, NoSectionError, NoOptionError
except:
    from configparser import ConfigParser, NoSectionError, NoOptionError

GRUB2_DIR = '/etc/default/grub.d'

#space reserved for swap headers
SWAP_RESERVED_SIZE = 16384
log_to_syslog = True
log_to_stderr = True
SWAP_FILE = '/swap-hibinit'
URL = "http://169.254.169.254/latest/meta-data/hibernation/configured"

def log(message):
    if log_to_syslog:
        syslog.syslog(message)
    if log_to_stderr:
        sys.stderr.write("%s\n" % message)

def sigterm_handler(signal, frame):
    #save the state here or do whatever you want
    log('Process killed cleaning up!')
    if os.path.isfile(SWAP_FILE) and os.access(SWAP_FILE, os.R_OK):
        os.remove(SWAP_FILE)
        exit(0)

def fallocate(fl, size):
    try:
        _libc = ctypes.CDLL('libc.so.6')
        _fallocate = _libc.fallocate
        _fallocate.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_ulong]

        # (FD, mode, offset, len)
        res = _fallocate(fl.fileno(), 0, 0, size)
        if res != 0:
            raise Exception("Failed to perform fallocate(). Result: %d" % res)
    except Exception as e:
        log("Failed to call fallocate(), will use resize. Err: %s" % str(e))
        fl.seek(size-1)
        fl.write(chr(0))


def get_file_block_number(filename):
    with open(filename, 'r') as handle:
        buf = array.array('L', [0])
        # from linux/fs.h
        FIBMAP = 0x01
        result = fcntl.ioctl(handle.fileno(), FIBMAP, buf)
    if result < 0:
        raise Exception("Failed to get the file offset. Error=%d" % result)
    return buf[0]

def get_rootfs_size():
    stat=os.statvfs('/')
    return math.ceil(float(stat.f_bsize * stat.f_blocks)/(1024*1024*1024))


def get_partuuid(device):
    return check_output(
        ['lsblk', '-dno', 'PARTUUID', device]).decode('ascii').strip()


def patch_grub_config(swap_device, offset, grub2_dir):
    """Update GRUB2 config when needed"""
    if grub2_dir and os.path.exists(grub2_dir):
        offset_file = os.path.join(grub2_dir, '99-set-swap.cfg')
        if swap_device.startswith("/dev"):
            swap_device = "PARTUUID=%s" % get_partuuid(swap_device)
        grub_snippet = (
            'GRUB_CMDLINE_LINUX_DEFAULT="$GRUB_CMDLINE_LINUX_DEFAULT'
            ' no_console_suspend=1 resume_offset=%d resume=%s"\n'
            % (offset, swap_device))
        if ((not os.path.exists(offset_file)
             or open(offset_file).read() not in (grub_snippet))):
            log("Updating GRUB to use the device %s with offset %d for resume"
                % (swap_device, offset))
            with open(offset_file, 'w') as fl:
                fl.write(grub_snippet)
            check_call('/usr/sbin/update-grub2')
            log("GRUB configuration is updated")


def update_kernel_swap_offset(swapon, swapoff, filename, grub_update):
    swapon = swapon.format(swapfile=filename)
    log("Running: %s" % swapon)
    check_call(swapon, shell=True)
    log("Updating the kernel offset for the swapfile: %s" % filename)

    statbuf = os.stat(filename)
    dev = statbuf.st_dev
    offset = get_file_block_number(filename)

    if grub_update:
        dev_str = find_device_for_file(filename)
        patch_grub_config(dev_str, offset, GRUB2_DIR)
    else:
        log("Skipping GRUB configuration update")

    log("Setting swap device to %d with offset %d" % (dev, offset))

    # Set the kernel swap offset, see https://www.kernel.org/doc/Documentation/power/userland-swsusp.txt
    # From linux/suspend_ioctls.h
    SNAPSHOT_SET_SWAP_AREA = 0x400C330D
    buf = struct.pack('LI', offset, dev)
    with open('/dev/snapshot', 'r') as snap:
        fcntl.ioctl(snap, SNAPSHOT_SET_SWAP_AREA, buf)
    log("Done updating the swap offset. Turning swapoff")
    swapoff = swapoff.format(swapfile=filename)
    log("Running: %s" % swapoff)
    check_call(swapoff, shell=True)


def find_device_for_file(filename):
    # Find the mount point for the swap file ('df -P /swap')
    df_out = check_output(['df', '-P', filename]).decode('ascii')
    dev_str = df_out.split("\n")[1].split()[0]
    return dev_str


class SwapInitializer(object):
    def __init__(self, filename, swap_size, touch_swap, mkswap, swapoff, swapon):
        self.filename = filename
        self.swap_size = swap_size
        self.mkswap = mkswap
        self.swapoff = swapoff
        self.swapon  = swapon
        self.touch_swap = touch_swap

    def do_allocate(self):
        log("Allocating %d bytes in %s" % (self.swap_size, self.filename))
        with open(self.filename, 'w+') as fl:
            fallocate(fl, self.swap_size)
            os.chmod(self.filename, 0o600)

    def init_swap(self):
        """
            Initialize the swap using direct IO to avoid polluting the page cache
        """
        try:
            cur_swap_size = os.stat(self.filename).st_size
            if cur_swap_size >= self.swap_size:
                log("Swap file size (%d bytes) is already large enough" % cur_swap_size)
                if self.init_mkswap():
                    return
        except OSError:
            try:
                os.unlink(self.filename)
            except:
                pass

        self.do_allocate()
        if not self.touch_swap:
            log("Swap pre-heating is skipped, the swap blocks won't be touched during "
                "to ensure they are ready")
            self.init_mkswap()
            return

        written = 0
        log("Opening %s for direct IO" % self.filename)
        fd = os.open(self.filename, os.O_RDWR | os.O_DIRECT | os.O_SYNC | os.O_DSYNC)
        if fd < 0:
            raise Exception("Failed to initialize the swap. Err: %s" % os.strerror(os.errno))

        filler_block = None
        try:
            # Create a filler block that is correctly aligned for direct IO
            filler_block = mmap.mmap(-1, 1024 * 1024)
            # We're using 'b' to avoid optimizations that might happen for zero-filled pages
            filler_block.write(b'b' * 1024 * 1024)

            log("Touching all blocks in %s" % self.filename)

            while written < self.swap_size:
                res = os.write(fd, filler_block)
                if res <= 0:
                    raise Exception("Failed to touch a block. Err: %s" % os.strerror(os.errno))
                written += res
        finally:
            os.close(fd)
            if filler_block:
                filler_block.close()
        log("Swap file %s is ready" % self.filename)
        self.init_mkswap()

    def init_mkswap(self):
        # Do mkswap
        try:
            mkswap = self.mkswap.format(swapfile=self.filename)
            log("Running: %s" % mkswap)
            check_call(mkswap, shell=True)
            return True
        except Exception as e:
            log("Failed to initialize swap, reason: %s" % str(e))
        return False


class BackgroundInitializerRunner(object):
    def __init__(self, swapper, update_grub):
        self.swapper = swapper
        self.thread = None
        self.error = None
        self.update_grub = update_grub
   
    def start_init(self):
        try:
            pid = os.fork()
            if pid > 0:
            # Exit parent process
                sys.exit(0)
        except OSError as e:
            print >> sys.stderr, "fork failed: %d (%s)" % (e.errno, e.strerror)
            sys.exit(1)

        # Configure the child processes environment
        os.chdir("/")
        os.setsid()
        os.umask(0o022)
  
    def do_async_init(self):
        try:
            self.swapper.init_swap()
            update_kernel_swap_offset(self.swapper.swapon, self.swapper.swapoff, self.swapper.filename, self.update_grub)
        except Exception as ex:
            log("Failed to initialize swap, reason: %s" % str(ex))
            self.error = ex


def swap_needs_touch(swapfile):
    # Walk the parent directories of the swapfile to find on which
    # filesystem it's mounted
    swap_place = swapfile
    dev = None
    while not dev:
        swap_place, _ = os.path.split(swap_place)        
        try:
            dev = find_device_for_file(swap_place)
        except:
            pass
            if swap_place == '/':
                raise Exception("Failed to find the filesystem type of /")

    with open("/proc/mounts") as fl:
        lines = fl.read().split("\n")
        for ln in lines:
            if dev in ln and "xfs" in ln:
                return True
    return False


class Config(object):
    def __init__(self, config, args):
        def get(section, name):
            try:
                return config.get(section, name)
            except NoSectionError:
                return None
            except NoOptionError:
                return None

        def get_int(section, name):
            v = get(section, name)
            if v is None:
                return None
            return int(v)

        self.log_to_syslog = self.merge(
            self.to_bool(get('core', 'log-to-syslog')), self.to_bool(args.log_to_syslog), True)
        self.log_to_stderr = self.merge(
            self.to_bool(get('core', 'log-to-stderr')), self.to_bool(args.log_to_stderr), True)

        self.mkswap = self.merge(get('swap', 'mkswap'), args.mkswap, 'mkswap {swapfile}')
        self.swapon = self.merge(get('swap', 'swapon'), args.swapon, 'swapon {swapfile}')
        self.swapoff = self.merge(get('swap', 'swapoff'), args.swapoff, 'swapoff {swapfile}')
        self.touch_swap = self.merge(
            self.to_bool(get('core', 'touch-swap')), self.to_bool(args.touch_swap),
            swap_needs_touch(SWAP_FILE))

        self.grub_update = self.merge(
            self.to_bool(get('core', 'grub-update')), self.to_bool(args.grub_update), True)

        self.swap_percentage = self.merge(
            get_int('swap', 'percentage-of-ram'), args.swap_ram_percentage, 100)
        self.swap_mb = self.merge(
            get_int('swap', 'target-size-mb'), args.swap_target_size_mb, 4000)


    def merge(self, cf_value, arg_value, def_val):
        if arg_value is not None:
            return arg_value
        if cf_value is not None:
            return cf_value
        return def_val

    def to_bool(self, bool_str):
        """Parse the string and return the boolean value encoded or raise an exception"""
        if bool_str is None:
            return None
        if bool_str.lower() in ['true', 't', '1']:
            return True
        elif bool_str.lower() in ['false', 'f', '0']:
            return False
        # if here we couldn't parse it
        raise ValueError("%s is not recognized as a boolean value" % bool_str)

    def __str__(self):
        return str(self.__dict__)

def hibernationEnabled():
     """Returns a boolean indicating whether hibernation is enabled or not."""
     response = None
     try:
         response = urlopen(URL)
         data = response.read()
         if data.lower() in ('false', b'false'):
             return False
     except:
         return False
     finally:
         if response:
             response.close()
     return True

def main():
    
    if not hibernationEnabled():
        log("Instance Launch has not enabled Hibernation Configured Flag. hibinit-agent exiting!!")
        exit(0)
    # Validate if disk space>total RAM 
    ram_bytes = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
    if get_rootfs_size()<=(math.ceil(float(ram_bytes)/(1024*1024*1024))):
        log("Insufficient disk space. Cannot create setup for hibernation. Please allocate a larger root device")
        exit(1)
    # Parse arguments
    parser = argparse.ArgumentParser(description="An EC2 background process that creates a setup for instance hibernation "
                                                 "at instance launch and also registers ACPI sleep event/actions")
    parser.add_argument('-c', '--config', help='Configuration file to use', type=str)
    parser.add_argument("-syslog", "--log-to-syslog", help='Log to syslog', type=str)
    parser.add_argument("-stderr", "--log-to-stderr", help='Log to stderr', type=str)
    parser.add_argument("-touch", "--touch-swap", help='Do swap initialization', type=str)
    parser.add_argument("-grub", "--grub-update", help='Update GRUB config with resume offset', type=str)
    parser.add_argument("-p", "--swap-ram-percentage", help='The target swap size as a percentage of RAM', type=int)
    parser.add_argument("-s", "--swap-target-size-mb", help='The target swap size in megabytes', type=int)
    parser.add_argument('--mkswap', help='The command line utility to set up swap', type=str)
    parser.add_argument('--swapon', help='The command line utility to turn on swap', type=str)
    parser.add_argument('--swapoff', help='The command line utility to turn off swap', type=str)

    args = parser.parse_args()

    config_file = ConfigParser()
    if args.config:
        config_file.read(args.config)

    config = Config(config_file, args)
    global log_to_syslog, log_to_stderr
    log_to_stderr = config.log_to_stderr
    log_to_syslog = config.log_to_syslog

    log("Effective config: %s" % config)

    target_swap_size = config.swap_mb * 1024 * 1024
    swap_percentage_size = ram_bytes * config.swap_percentage // 100
    if swap_percentage_size > target_swap_size:
        target_swap_size = int(swap_percentage_size)
    log("Will check if swap is at least: %d megabytes" % (target_swap_size // (1024*1024)))

    #Validate if swap file exists
    cur_swap = 0 
    if os.path.isfile(SWAP_FILE) and os.access(SWAP_FILE, os.R_OK):
        cur_swap = os.path.getsize(SWAP_FILE)

    bi = None
    if cur_swap >= target_swap_size - SWAP_RESERVED_SIZE:
        log("There's sufficient swap available (have %d, need %d)" %
            (cur_swap, target_swap_size))
        update_kernel_swap_offset(config.swapon, config.swapoff, SWAP_FILE, config.grub_update)
        exit()
    #validate if instance was launched from pre-created image and swap size>=total RAM, if not re-create the swap 
    elif cur_swap > 0 and (cur_swap < target_swap_size - SWAP_RESERVED_SIZE):
        log("Swap already exists! (have %d, need %d), deleting existing swap file %s" %
		(cur_swap, target_swap_size, SWAP_FILE))
        os.remove(SWAP_FILE)
        
    log("Create swap and initialize it")
    # We need to create swap, but first validate that we have enough free space
    swap_dev = os.path.dirname(SWAP_FILE)
    st = os.statvfs(swap_dev)
    free_bytes = st.f_bavail * st.f_frsize
    #rounding off to swap_size+10mb for swap headers
    free_space_needed = target_swap_size + 10 * 1024 * 1024
    if free_space_needed >= free_bytes:
        log("There's not enough space (%d present, %d needed) on the device: %s" % (
		free_bytes, free_space_needed, swap_dev))
        exit(1)
        log("There's enough space (%d present, %d needed) on the device: %s" % (
            free_bytes, free_space_needed, swap_dev))

    sw = SwapInitializer(SWAP_FILE, target_swap_size, config.touch_swap,
                             config.mkswap, config.swapoff, config.swapon)
    bi = BackgroundInitializerRunner(sw, config.grub_update)
    signal.signal(signal.SIGTERM, sigterm_handler)

    if bi:
        bi.start_init()
        
    log("kicking child process to initiate the setup")
    bi.do_async_init()

if __name__ == '__main__':
    main()