#!/usr/bin/python3

import argparse
import ctypes
import hashlib
import os
import re
import shutil
import subprocess
import sys
import tempfile
import uuid

_NULL = open('/dev/null', 'wb')

if os.geteuid() != 0:
    print("Super-user privileges are required to reflash endpoints")
    sys.exit(1)
    
_REFLASH_PATH = "/boot/iep-flash.img"
_UPDATE_PATH = "/var/3d-p/firmware/update"
_OPTIONS_PATH = "/var/3d-p/firmware/update-options"
_VERSION_PATH = "/etc/3d-p/firmware/version"
_REFLASH_HELPER = "/usr/share/iep-maintenance/scripts/reflash-helper"

parser = argparse.ArgumentParser(description="Reinstall the Intelligent Endpoint operating system")
parser.add_argument('image_file', type=str, default=_REFLASH_PATH, nargs='?', help='The image-file to be flashed')
parser.add_argument('--timeout', type=int, default=300, help="The number of seconds to wait for interaction during an attended reflash before aborting")
parser.add_argument('--no-preserve', dest="preserve", action="store_false", help="Treats the reflash as a complete config-wipe, suppressing the usual attempts at maintaining post-reflash connectivity")
parser.add_argument('--unattended', action='store_true', default=False, help="Whether default values should be assumed to allow the reflash to proceed without input")
parser.add_argument('--cancel', action='store_true', default=False, help="Cancel a previously staged reflash")
args = parser.parse_args()

if args.cancel:
    try:
        os.unlink(_UPDATE_PATH)
    except OSError as e:
        if e.errno == 2:
            print("No reflash is staged")
        else:
            raise
    else:
        print("Reflash cancelled")
    sys.exit(0)
    
if args.image_file != _REFLASH_PATH:
    if os.getenv('SUDO_UID') is not None:
        print("New flash images may not be installed via sudo")
        sys.exit(1)
        
def checksum(flash_image):
    if not os.path.isfile(flash_image):
        raise IOError("{file} does not appear to exist".format(
            file=flash_image,
        ))
        
    try:
        subprocess.check_call(('checkisomd5', flash_image,), stdout=_NULL, stderr=_NULL)
    except Exception as e:
        raise IOError("Unable to validate the integrity of {file}".format(
            file=flash_image,
        ))
        
_boot_mounted = False
_boot_writeable = False
try: #See if /boot is mounted
    subprocess.check_call(('/bin/mountpoint', '-q', '/boot'))
except Exception as e: #Mount it
    subprocess.check_call(('/bin/mount', '-o', 'rw', '/boot'))
else:
    _boot_mounted = True
    _boot_writeable = any(
        l for l in subprocess.check_output(('/bin/mount',)).splitlines()
                if b'on /boot ' in l and b'rw' in l
    )
    if not _boot_writeable:
        subprocess.check_call(('/bin/mount', '-o', 'remount,rw', '/boot'))
        
try:
    print("Extracting data from the image...")
    image_version = None
    image_timestamp = None
    iso_dir = tempfile.mkdtemp()
    try:
        subprocess.check_call(('/bin/mount', '-o', 'ro', args.image_file, iso_dir))
        try:
            if os.path.isfile(os.path.join(iso_dir, 'type')):
                if open(os.path.join(iso_dir, 'type'), 'rb').read().strip() != b'flash:ties':
                    raise ValueError("The supplied archive is not a Ties flash image")
            elif not os.path.isfile(os.path.join(iso_dir, 'archives', 'firmware.tar.xz')):
                raise ValueError("The supplied archive does not appear to be a flash image")
                
            running_bom = open("/boot/3d-p/bom", 'r').read().strip()
            for bom in os.listdir(os.path.join(iso_dir, 'boms')):
                if re.match('^{}$'.format(bom.replace('_', '.')), running_bom):
                    break
            else:
                raise ValueError("The given flash image does not support this BoM")
                
            image_version = open(os.path.join(iso_dir, 'version'), 'r').read().strip()
            if os.path.isfile(os.path.join(iso_dir, 'timestamp')):
                image_timestamp = open(os.path.join(iso_dir, 'timestamp'), 'r').read().strip()
        finally:
            subprocess.check_call(('/bin/umount', iso_dir))
    finally:
        shutil.rmtree(iso_dir)
        
    if args.image_file != _REFLASH_PATH:
        print("Verifying integrity of {path}...".format(
            path=args.image_file,
        ))
        checksum(args.image_file)
        
        reflash_size = os.stat(args.image_file).st_size
        boot_stat = os.statvfs('/boot')
        boot_free_space = boot_stat.f_bavail * boot_stat.f_bsize 
        boot_file_exists = os.path.isfile(_REFLASH_PATH)
        if boot_file_exists:
            boot_reflash_size = os.stat(_REFLASH_PATH).st_size
        else:
            boot_reflash_size = 0
            
        if reflash_size <= boot_free_space:
            flash_tempfile = '/boot/' + str(uuid.uuid1())
            print("Moving new reflash file to {path}...".format(path=flash_tempfile))
            shutil.move(args.image_file, flash_tempfile)
            print("Verifying integrity of reflash file again for paranoia's sake...")
            try:
                checksum(flash_tempfile)
            except ValueError:
                os.unlink(flash_tempfile)
                raise
                
            #Point of no return
            print("Moving new reflash file to {path}...".format(path=_REFLASH_PATH))
            shutil.move(flash_tempfile, _REFLASH_PATH)
        elif reflash_size <= boot_free_space + boot_reflash_size:
            #No way around it; trust that things will succeed
            print("Removing old reflash file at {path} to free space...".format(path=_REFLASH_PATH))
            os.unlink(_REFLASH_PATH)
            print("Moving new reflash file to {path}...".format(path=_REFLASH_PATH))
            shutil.move(args.image_file, _REFLASH_PATH)
        else:
            raise IOError("Insufficient space available for the reflash operation")
            
    if not os.path.isfile(_REFLASH_PATH):
        raise EnvironmentError("No reflash file found")
        
    #Compare the new version to the installed version to inform the user, if appropriate
    installed_version = open(_VERSION_PATH, 'r').read().strip()
    print("Installed version: {installed}".format(
        installed=installed_version,
    ))
    print("Version to be installed: {image} ({timestamp})".format(
        image=image_version,
        timestamp=(image_timestamp or "build time unavailable"),
    ))
    if installed_version > image_version:
        print('\033[5mYOU ARE ABOUT TO DOWNGRADE TO AN OLDER FIRMWARE VERSION\033[0m')
finally:
    ctypes.CDLL("libc.so.6").sync()
    if not _boot_mounted:
        subprocess.check_call(('/bin/umount', '/boot'))
    elif not _boot_writeable:
        subprocess.check_call(('/bin/mount', '-o', 'remount,ro', '/boot'))
        
if os.path.isfile(_UPDATE_PATH):
    print("Another update was previously slated; its instructions have been relocated to {path}.old".format(path=_UPDATE_PATH))
    shutil.move(_UPDATE_PATH, _UPDATE_PATH + '.old')
    
print("")
options = {
    'UPDATE_OPTIONS_TIMEOUT': args.timeout,
    'UPDATE_OPTIONS_UNATTENDED': args.unattended and 'y' or 'n',
    'UPDATE_OPTIONS_PRESERVE_CONFIG': args.preserve and 'y' or 'n',
}
if args.unattended:
    print("#" * 80)
    print('{:^80}'.format("Unattended remote reflashes are NOT RECOMMENDED by 3D-P"))
    print("#" * 80)
    options['UPDATE_OPTIONS_TIMEOUT'] = 0
    
if args.preserve:
    print("#" * 80)
    print('{:^80}'.format("Connectivity configuration will be preserved if possible"))
    print('{:^80}'.format("If a complete reset is desired, pass --no-preserve"))
    print("#" * 80)
    
open(_OPTIONS_PATH, 'w').writelines("{k}={v}\n".format(k=k, v=v) for (k, v) in options.items())
shutil.copy2(_REFLASH_HELPER, _UPDATE_PATH)

print("The reflash operation will begin when you reboot the Intelligent Endpoint")
print("Remote reflashes are not recommended; if possible, proceed only with a serial connection and a stable power source")
print("To begin, 'reboot'")
print("To cancel, '{} --cancel'".format(sys.argv[0]))
