#!/usr/bin/python

"""
Written by Igor Pashev <pashev.igor@gmail.com>

The author has placed this work in the Public Domain,
thereby relinquishing all copyrights. Everyone is free
to use, modify, republish, sell or give away this work
without prior consent from anybody.
"""

from lib.hdd import HDD
from lib.snack import *
from snack import *
from subprocess import Popen, PIPE, call
from pprint import pprint
from tempfile import mkstemp
from urllib2 import urlopen, URLError, HTTPError
import os
import re
import sys

# Snack screen
screen = None

# List of disks installed on the system:
hdds = []

# Name of the root ZFS pool: either existing or newly created
rpool = None

# Name of the root slice, where ZFS root pool will be created (e. g. c0t0d0s0)
rslice = None

# Created boot environment
bootenv = None

# Version to install:
codename = 'bok'

# Dyson mirrors:
mirror = None
mirrors = [
    ('http://apt.osdyson.org', 'Russia'),
    ('http://mirror-us.osdyson.org/apt/', 'USA'),
    ]


# FS to be created in boot environment.
# Order is important.
befs = [
    {'name':'/usr', 'options':{}},
    {'name':'/usr/local', 'options':{
                                'compression':'on',
                                }},
    {'name':'/usr/src', 'options':{
                                'compression':'on',
                                }},
    {'name':'/var', 'options':{}},
    {'name':'/var/cache', 'options':{}},
    {'name':'/var/lib', 'options':{}},
    {'name':'/var/lib/dpkg', 'options':{
                                'compression' : 'on',
                                'setuid'      : 'off',
                                'devices'     : 'off',
                                }},
    {'name':'/var/log', 'options':{
                                'compression' : 'on',
                                'setuid'      : 'off',
                                'devices'     : 'off',
                                }},
    {'name':'/var/mail', 'options':{
                                'compression' : 'on',
                                'setuid'      : 'off',
                                'devices'     : 'off',
                                }},
    {'name':'/var/spool', 'options':{
                                'compression' : 'on',
                                'setuid'      : 'off',
                                'devices'     : 'off',
                                }},
    {'name':'/var/tmp', 'options':{}},
    ]

class Abort(Exception):
    ''' User clicked "Cancel" '''
    pass
class NoDisks(Exception):
    ''' No Hard Drives found '''
    pass

def welcome():
    welcome = ButtonChoiceWindow(screen, "Welcome to the Dyson installer",
            '''
Dyson is an operating system derived from Debian and based on Illumos core. \
It uses Illumos unix kernel and libc, ZFS file system and SMF to manage system startup. \
To learn more visit http://osdyson.org

This installer will guide you through disk paritioning, filesystem creation, \
installation of base system and minimal configuration.

Please note, this system IS VERY EXPERIMENTAL. \
You can LOOSE ALL YOUR DATA which this system can reach ;-)

Whould you like to continue?''',
            buttons=[('Install Dyson', True), ('Exit the installer', False)], width=70)
    if not welcome:
        raise Abort('Installation canceled')

def find_zpools():
    global zpools
    progress = ProgressMessage(screen, title='Please, wait', text='Searching for ZFS pools ...')
    zpools = []
    zpool_cmd = Popen(['zpool', 'list', '-H', '-o', 'name'], stdout=PIPE, stderr=PIPE)
    for line in zpool_cmd.stdout:
        zpools.append(line.rstrip())
    zpool_cmd = Popen("zpool import | awk '/pool:/ {print $2}'", shell=True, stdout=PIPE, stderr=PIPE)
    for line in zpool_cmd.stdout:
        zpools.append(line.rstrip())

def find_hdds():
    global hdds
    progress = ProgressMessage(screen, title='Please, wait', text='Searching for disks ...')
    hdds = []
    pat = re.compile('\d+\.\s+(\S+)\s+<(.+?)(\s+cyl.*)?>')
    format_cmd = Popen('format </dev/null', shell=True, stdout=PIPE)
    for line in format_cmd.stdout:
        m = pat.search(line)
        if m:
            name = m.group(1)
            desc = m.group(2)
            hdds.append(HDD(name, desc))
    if len(hdds) == 0:
        raise NoDisks('No disks found')


def choose_rpool():
    zpool_items = map(lambda x: (x, x), zpools)
    zpool_items.append(('Create new ZFS pool', None))

    choice = None

    buttons = [('Ok', 'ok'), ('Cancel', 'cancel')]

    while not choice in ['ok', 'cancel']:
        (choice, pool) = ListboxChoiceWindow(screen,
                title='Choose ZFS pool',
                text='Choose ZFS pool where Dyson will be installed.\n'
                'You can selected one of the existing pools or create a new one.',
                items=zpool_items,
                buttons=buttons,
                width=76,
                scroll=1,
                default=0,
                )
        # Enter is pressed:
        if choice == None:
            choice = 'ok'
    if choice == 'ok':
        return pool
    raise Abort('Pool choosing')

def create_root_slice(hdd_name):
    '''Create root slice (0) over entire Solaris partition'''
    progress = ProgressMessage(screen, title='Please, wait',
            text='Formatting Solaris partition on disk "{}" ...'.format(hdd_name))

    fd, path = mkstemp()
    os.write(fd, 'partition\nmodify\n1\nyes\n0\n\n\n\n\n\n\nyes\n"{label}"\nquit\nquit\n'.format(label=hdd_name))
    os.close(fd)
    format_cmd = Popen(['format', '-f', path, '-d', hdd_name], stderr=PIPE, stdout=PIPE)
    out, err = format_cmd.communicate()
    os.unlink(path)
    if format_cmd.returncode != 0:
        ButtonChoiceWindow(screen, title='Creating of root slice failed',
                text=err, buttons = ['Ok'], width=70)
        return None
    return '{}s0'.format(hdd_name)

def create_rpool(hdd):
    global rslice
    rslice = create_root_slice(hdd.name)
    if not rslice:
        return None
    pool_name = None
    for pool_no in ['', 1, 2, 3, 4, 5, 6, 7]:
        pool_name = 'rpool{}'.format(pool_no)
        if not pool_name in zpools:
            break
        else:
            pool_name = None;

    progress = ProgressMessage(screen, title='Please, wait',
            text='Creating ZFS pool "{rpool}" on slice "{slice}" ...'.format(rpool=pool_name, slice=rslice))

    zpool_cmd = Popen(['zpool', 'create', '-R', '/mnt/'+pool_name, '-f', pool_name, rslice],
            stderr=PIPE, stdout=PIPE)
    out, err = zpool_cmd.communicate()
    if 0 != zpool_cmd.returncode:
        ButtonChoiceWindow(screen,
                title='Creating of ZFS root pool failed',
                text=err, buttons = ['Ok'], width=70)
        return None
    return pool_name

def pool_is_imported(pool):
    rpool_is_imported = call(['zpool', 'status', pool], stdout=PIPE, stderr=PIPE)
    return rpool_is_imported == 0

def configure_rpool():
    global rpool
    find_zpools()
    while not rpool:
        if len(zpools) > 0:
            rpool = choose_rpool()
        if not rpool:
            hdd = choose_hdd(len(zpools))
            if hdd:
                rpool = create_rpool(hdd)
        if rpool:
            if not pool_is_imported(rpool):
                progress = ProgressMessage(screen, title='Please, wait',
                        text='Importing ZFS pool "{}" ...'.format(rpool))
                zpool_cmd = Popen(['zpool', 'import', '-R', '/mnt/'+rpool, '-f', rpool], stderr=PIPE, stdout=PIPE)
                out, err = zpool_cmd.communicate()
                if 0 !=  zpool_cmd.returncode:
                    ButtonChoiceWindow(screen,
                            title='Importing of ZFS pool "{}" failed'.format(rpool),
                            text=err, buttons=['Ok'], width=70)
                    rpool = None

def zfs_exists(fs):
    fs_exists = call(['zfs', 'list', '-H', fs], stdout=PIPE, stderr=PIPE)
    return fs_exists == 0

def zfs_create(fs, zvol=None, options={}):
    args = ['zfs', 'create']
    if zvol:
        args += ['-V', zvol]

    for o in options:
        args += ['-o', o + '=' + options[o]]

    args.append(fs)

    zfs_cmd = Popen(args, stderr=PIPE, stdout=PIPE)
    out, err = zfs_cmd.communicate()
    if 0 != zfs_cmd.returncode:
        ButtonChoiceWindow(screen,
                title='Creating of ZFS dataset failed',
                text=err, buttons = ['Ok'], width=70)
        return False
    return True


def configure_zfs():
    global bootenv
    progress = ProgressBar(screen,
            title='Creating ZFS filesystems',
            top = 4 + len(befs), # 4 for ROOT, /swap, /home, and BE
            width = 70,
            )
    root = rpool + '/ROOT'
    if not zfs_exists(root):
        progress.text = 'Creating ' + root
        zfs_create(root, options={'canmount':'off'})
    progress.advance()

    swap = rpool + '/swap'
    if not zfs_exists(swap):
        progress.text = 'Creating ' + swap
        zfs_create(swap, zvol='2G')
    progress.advance()

    home = rpool + '/home'
    if not zfs_exists(home):
        progress.text = 'Creating ' + home
        zfs_create(home, options={'mountpoint':'/home', 'compression':'on',
            'devices':'off', 'setuid':'off'})
    progress.advance()

    # Make sure we are creating new boot environment
    bootenv = None
    for be_no in range(20):
        bootenv = root + '/osdy-' + str(be_no)
        if not zfs_exists(bootenv):
            break
        else:
            bootenv = None

    progress.text = 'Creating ' + bootenv
    zfs_create(bootenv, options={'canmount':'noauto'})
    progress.advance()

    for fs in befs:
        p = bootenv + fs['name']
        progress.text = 'Creating ' + p
        zfs_create(p, options=fs['options'])
        progress.advance()

def choose_hdd(number_of_zpools = 0):
    hdd_items = map(lambda x: '{}: {} {}'.format(x.name, x.capacity, x.description), hdds)
    choice = None

    buttons = []
    buttons.append(('Use selected disk', 'ok'))
    if number_of_zpools > 0:
        buttons.append(('Back', 'zpool'))
    buttons.append(('Cancel', 'cancel'))


    while not choice in ['ok', 'zpool', 'cancel']:
        (choice, hdd_no) = ListboxChoiceWindow(screen,
                title='Choose hard disk drive',
                text='Choose a disk for the root ZFS pool',
                items=hdd_items,
                buttons=buttons,
                width=76,
                scroll=1,
                default=0,
                help=None
                )
        if choice in ['ok', None]:
            choice = configure_partitions(hdds[hdd_no], len(hdds))
            if choice == 'another':
                choice = None

    if choice == 'ok':
        return hdds[hdd_no]
    if choice == 'zpool':
        return None
    raise Abort('HDD choosing')

def configure_partitions(hdd, number_of_disks = 1):
    choice = None
    top_text = 'For the root ZFS pool you need a disk containing a Solaris partition (id 0xbf). \
If this disk includes such a partition and you are happy with it, use this disk. \
Otherwise you have to change partitioning. Note that only the first Solaris partition \
will be used.'

    while not choice in ['ok', 'another', 'cancel']:
        part_info = top_text
        part_info += '\n\n'

        have_partitions = len(hdd.partitions) > 0
        have_solaris_partition = False
        have_gpt = False
        buttons = []

        if have_partitions:
            part_info += 'Partitions on {name} {desc} {size}:\n\n'.format(
                    name=hdd.name, desc=hdd.description, size=hdd.capacity)
            part_info += '    TYPE    ID    NAME                     SIZE\n'
            n = 0
            for p in hdd.partitions:
                n = n + 1
                if p.id == 0xBF:
                    have_solaris_partition = True
                if p.id == 0xEE:
                    have_gpt = True
                part_info += '#{n:<2} {primary} {id:#04x} {system:25} {capacity:10}\n'.format(
                    n=n, capacity=p.capacity, system=p.system, id=p.id,
                    primary=['primary', 'logical'][p.logical])
        else:
            part_info += 'This disk is not partitioned.\n'

        if have_solaris_partition:
            buttons.append(('Use this disk', 'ok'))

        if have_gpt:
            part_info += '\nThis disk has GUID Partition Table (GPT) and is unsupported.\n'
            # XXX gparted
        else:
            buttons.append(('Edit partitions', 'fdisk'))

        if number_of_disks > 1:
            buttons.append(('Back', 'another'))

        buttons.append(('Cancel', 'cancel'))

        choice = ButtonChoiceWindow(screen,
                title='Partitions of {}'.format(hdd.name),
                text=part_info,
                buttons=buttons,
                width=76,
                help=None
                )
        if choice == 'fdisk':
            screen.suspend()
            call(['cfdisk', hdd.raw_device])
            hdd.reread_partitions()
            screen.resume()
    # end while

    if choice == 'cancel':
        raise Abort('Partitioning')
    return choice

def get_mirror():
    if not hasattr(get_mirror, "m"):
        get_mirror.m=''
    while True:
        (choice, entry) = EntryWindow(screen,
                title='Enter APT repository URL',
                text='If you want to have line '
                '"deb ftp://192.168.1.1 bok main" in /etc/apt/sources.list '
                'enter "ftp://192.168.1.1" here.\n'
                'Remember to add URL protocol.\n',
                prompts=[('URL', get_mirror.m)],
                buttons = [ ('Ok', 'ok'), ('Back', 'back'), ('Cancel', 'cancel')],
                width=70,
                entryWidth=50,
                )
        if choice == 'back':
            return None
        if choice in ['ok', None]:
            get_mirror.m = entry[0].strip()
            if valid_mirror(get_mirror.m):
                return get_mirror.m
        if choice == 'cancel':
            raise Abort('Entering an APT mirror')


def valid_mirror(mirror):
    progress = ProgressMessage(screen, title='Please, wait',
            text='Checking APT mirror: ' + mirror)
    try:
        o = urlopen('{mirror}/dists/{codename}/Release'.format(
            mirror=mirror, codename=codename), timeout=10)
        for line in o.read(400).split('\n'):
            if line.strip() == 'Codename: ' + codename:
                return True
        raise URLError('Not an APT repository')
    except HTTPError as e:
        ButtonChoiceWindow(screen,
                title='Error',
                text='{code} {reason}.\n\nMirror {mirror} is not usable.'.format(
                    mirror=mirror, code=e.code, reason=e.reason),
                buttons = ['Ok'], width=70)
    except URLError as e:
        ButtonChoiceWindow(screen,
                title='Error',
                text='{reason}.\n\nMirror {mirror} is not usable.'.format(
                    reason=e.reason,mirror=mirror),
                buttons = ['Ok'], width=70)
    except ValueError:
        ButtonChoiceWindow(screen,
                title='Failed to check mirror',
                text='Invalid URL: ' + mirror,
                buttons = ['Ok'], width=70)
    except Exception:
        ButtonChoiceWindow(screen,
                title='Failed to check mirror',
                text='Unknown error: ' + mirror,
                buttons = ['Ok'], width=70)
    return False

def configure_mirror():
    global mirror
    items = map(lambda x: '{url} ({info})'.format(url=x[0], info=x[1]), mirrors)
    items.append(('Enter another  mirror', None))
    while not mirror:
        (choice, m) = ListboxChoiceWindow(screen,
                title='Choose APT mirror',
                text='Choose APT repository from which Dyson will be installed',
                items=items,
                buttons=[('Ok','ok'), ('Cancel', 'cancel')],
                width=76,
                scroll=1,
                default=0,
                )
        if choice in ['ok', None]:
            if m == None:
                mirror = get_mirror()
            else:
                if valid_mirror(mirrors[m][0]):
                    mirror = mirrors[m][0]
        if choice == 'cancel':
            raise Abort('Choosing APT mirror')


def debootstrap():
    progress = ProgressBar(screen,
            title='Installing base system',
            width=70)

    progress.text = 'Starting debootstrap ...'

    read, write = os.pipe()

    pid = os.fork()
    if pid == 0:
        os.close(read)
        os.dup2(write, 3)
        os.close(write)
        try:
            os.execl('/usr/sbin/debootstrap', 'debootstrap',
                    '--debian-installer', '--no-check-gpg',
                    '--include=illumos-grub,illumos-kernel',
                    codename, '/mnt/' + rpool + '/' + bootenv, mirror)
        except OSError as e:
            sys.exit(e.errno)
    else:
        os.close(write)
        f = os.fdopen(read, 'r', 0)
        r = re.compile(' +')
        while True:
            line = f.readline()
            if line:
                progress.text = line
            else:
                break



def destroy_bootenv():
    pass

screen = SnackScreen()
screen.pushHelpLine(' ')
try:
    welcome()
    find_hdds()
    configure_rpool()
    configure_mirror()
    configure_zfs()
    debootstrap()

except Abort as e:
    destroy_bootenv()
    pass
except NoDisks as e:
    print (e)
finally:
    screen.finish()

sys.exit(0)