#!/usr/bin/python
### BEGIN INIT INFO
# Provides: slapos_firstboot
# Required-Start:
# Required-Stop:
# Default-Start:
# Default-Stop:
# Description: Configures first boot of SlapOS
### END INIT INFO
##############################################################################
#
# Copyright (c) 2010 Vifib SARL and Contributors. All Rights Reserved.
#
# WARNING: This program as such is intended to be used by professional
# programmers who take the whole responsibility of assessing all potential
# consequences resulting from its eventual inadequacies and bugs
# End users who are looking for a ready-to-use solution with commercial
# guarantees and support are strongly advised to contract a Free Software
# Service Company
#
# This program is Free Software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
##############################################################################
import os
import shutil
import subprocess
import tempfile
import time

# some constants
LABEL="SLAPOS"
MINIMUM_FREE_SPACE_GB=60
MINIMUM_FREE_SPACE=MINIMUM_FREE_SPACE_GB * 1024 * 1024 * 1024
SLAPOS_MARK='# Added by SlapOS\n'

def callWithCheck(args):
  """Calls args in subprocesses, rasies ValueError if command failed"""
  if subprocess.call(args) != 0:
    raise ValueError

def callWithIgnore(args):
  """Calls args in subprocess, ignores issues"""
  subprocess.call(args)

def partprobe():
  """Calls partprobe utility"""
  callWithIgnore("partprobe")

def env():
  """Returns language neutreal environment"""
  env = {}
  for k, v in os.environ.iteritems():
    if 'LANG' not in k:
      env[k] = v
  env['LANG'] = 'C'
  env['LANGUAGE'] = 'C'
  return env

def getPartitionList(path):
  """Returns list of all partitions found on disk"""
  partprobe()
  result = subprocess.Popen(['fdisk', '-l', path], env=env(),
        stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()[0]
  partition_list = []
  for l in result.split('\n'):
    if l.startswith(path):
      partition_list.append(l.split()[0])
  return partition_list

def setupFilesystem(path):
  """Setups partition on path"""
  print "Setting up filesystem on %r" % path
  partprobe()
  callWithCheck(["mkfs", "-m", "1", "-L", LABEL, "-t", "ext4", path])
  partprobe()

def prepareDisk(path):
  """Prepares disk and creates partition"""
  callWithCheck(["parted", "--script", path, "mklabel", "msdos"])
  callWithCheck(["parted", "--script", '--', path, "mkpart", "primary",
    "ext2", "0", "-1"])
  partition = path + '1'
  setupFilesystem(partition)

def findEmptyDisk():
  """Tries to find not configured yet disk"""
  result = subprocess.Popen(['parted', '--script', '-l'], env=env(),
      stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate()[0]
  for line in result.split('\n'):
    if 'unrecognised disk label' in line:
      disk = line.split(':')[1].strip()
      fdisk = subprocess.Popen(['fdisk', '-l', disk], env=env(),
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate()[0]
      size = 0
      for l in fdisk.split('\n'):
        if l.startswith('Disk %s:' % disk):
          size = int(l.split(' ')[4])
          break
      if size > MINIMUM_FREE_SPACE:
        return disk

def deactivateSwap():
  """Disactivates swap in /etc/fstab and turns off any existing"""
  old_fstab = open('/etc/fstab', 'r').readlines()
  fstab = open('/etc/fstab', 'w')
  for line in old_fstab:
    if 'swap' in line:
      line = '#deactivated by slapos#' + line
    fstab.write(line)
  fstab.close()
  callWithIgnore(['swapoff', '-a'])

def prepareSlapOSPartition(mount_point):
  """Prepare SlapOS partitioning"""
  new_fstab = open('/etc/fstab', 'r').readlines()
  a = new_fstab.append
  d = dict(label=LABEL)
  a(SLAPOS_MARK)
  a('LABEL=%(label)s /mnt/%(label)s ext4 defaults,noatime 0 0\n'%
      d)
  a("/mnt/%(label)s/opt /opt none bind 0 0\n" % d)
  a("/mnt/%(label)s/tmp /tmp none bind 0 0\n" % d)
  a("/mnt/%(label)s/srv /srv none bind 0 0\n" % d)
  open('/etc/fstab', 'w').write(''.join(new_fstab))
  partprobe()
  retry = 10
  while True:
    try:
      print 'Trying to mount new partition, %s to go.' % retry
      callWithCheck(['mount', '/mnt/'+LABEL])
    except ValueError:
      if retry > 0:
        retry -= 1
        # give some time
        time.sleep(2)
      else:
        raise
    else:
      break
  for d in ['opt', 'srv', 'tmp']:
    p = '/mnt/'+LABEL+'/'+d
    if not os.path.exists(p):
      os.mkdir(p)
      os.chmod(p, 0755)
  os.chmod('/mnt/'+LABEL+'/tmp', 01777)
  callWithCheck(['mount', '/opt'])
  callWithCheck(['mount', '/tmp'])
  callWithCheck(['mount', '/srv'])

def configureNtp():
  """Configures NTP daemon"""
  server = "server pool.ntp.org"
  old_ntp = open('/etc/ntp.conf', 'r').readlines()
  new_ntp = open('/etc/ntp.conf', 'w')
  for line in old_ntp:
    if line.startswith('server'):
      continue
    new_ntp.write(line)
  new_ntp.write(SLAPOS_MARK)
  new_ntp.write(server+'\n')
  new_ntp.close()
  callWithIgnore(['chkconfig', '--add', 'ntp'])

def getMountedPartitionList():
  partition_list = []
  for line in open('/etc/mtab', 'r').readlines():
    if line.startswith('/dev'):
      partition_list.append(line.split(' ')[0])
  return partition_list

def getWidestFreeSpaceList(disk_dict):
  free_space_disk_dict = {}
  for disk in disk_dict.iterkeys():
    v = disk_dict[disk]
    for l in v:
      size = int(l.split(':')[3].rstrip('B'))
      free_space_disk_dict.setdefault(size, [])
      free_space_disk_dict[size].append(dict(disk=disk, info=l))
  if free_space_disk_dict:
    m = max(free_space_disk_dict)
    if m > MINIMUM_FREE_SPACE:
      return free_space_disk_dict[max(free_space_disk_dict)]
  return []

def findAndPrepareFreeSpace():
  """Finds free space, ignoring / mounted device
  Prefers disk with correct empty partition table"""
  mounted_partition_list = getMountedPartitionList()
  disk_dict = {}
  parted = subprocess.Popen(['parted', '--script', '-lm'], env=env(),
      stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()[0]
  for line in parted.split('\n'):
    line = line.strip()
    if line.startswith('/dev'):
      disk = line.split(':')[0]
      if any([q.startswith(disk) for q in mounted_partition_list]):
        continue
      # disk found, time to fetch free space
      disk_dict[disk] = []
      free = subprocess.Popen(['parted', '--script', '-m', disk, 'unit', 'B', 'print',
        'free'], env=env(),
        stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()[0]
      for f in free.split('\n'):
        if ':free;' in f:
          disk_dict[disk].append(f.strip())

  disk_space_list = getWidestFreeSpaceList(disk_dict)
  if len(disk_space_list) == 0:
    raise ValueError('Minumum of free space %sGB not found' % MINIMUM_FREE_SPACE_GB)
  disk = disk_space_list[0]
  beg = disk['info'].split(':')[1].rstrip('B')
  end = disk['info'].split(':')[2].rstrip('B')
  before = getPartitionList(disk['disk'])
  callWithCheck(['parted', '--script', '-m', '-a', 'optimal', '--',
    disk['disk'], 'unit', 'B', 'mkpart', 'primary', 'ext2', beg, end])
  after = getPartitionList(disk['disk'])
  new = [q for q in after if q not in before]
  setupFilesystem(new[0])
  return disk['disk']

def configureGrub(mount_point, slapos_label_file):
  """Configures grub on mount_point using disk"""
  boot = os.path.join(mount_point, 'boot')
  grub = os.path.join(boot, 'grub')
  if not os.path.exists(boot):
    os.mkdir(boot)
  if os.path.exists(grub):
    shutil.rmtree(grub)
  shutil.copytree('/boot/grub', grub)
  open(os.path.join(grub, 'menu.lst'), 'w').write("""timeout 3600
title SlapOS Error: The USB key is not first and bootable device
root (hd0,1)
chainloader +1
title Solution: Put the USB key
root (hd0,1)
chainloader +1
title Solution: Configure BIOS -- boot from USB key
root (hd0,1)
chainloader +1
title Solution: Configure BIOS -- USB key as first device
root (hd0,1)
chainloader +1
""")

  new_map_path = tempfile.mkstemp()[1]
  new_map = open(new_map_path, 'w')
  disk = os.path.realpath(os.path.join(os.path.dirname(slapos_label_file),
      os.readlink(slapos_label_file)))
  # this represents /dev/sdXN, so simply let's remove all possible numbers
  # from right
  # Note: This method's perfectioness is same as rest in this script, high
  # dependency on system configuration
  disk = disk.rstrip('0123456789')
  new_map.write(open('/boot/grub/device.map').read() + '\n(hd1) %s'% disk)
  new_map.close()
  grub_install = subprocess.Popen(['grub', '--batch', '--device-map=%s' %
    new_map], stdin=subprocess.PIPE)
  grub_install.communicate("""root (hd1,0)
setup (hd1)
quit
""")
  os.unlink(new_map_path)
  if grub_install.returncode is None or grub_install.returncode != 0:
    raise ValueError('Issue during grub installation')

def run():
  """Prepares machine to run SlapOS"""
  print "Running SUSE Studio first boot script..."
  partprobe()
  slapos_label_file='/dev/disk/by-label/' + LABEL
  if not os.path.exists(slapos_label_file):
    empty_disk = findEmptyDisk()
    if empty_disk is not None:
      print "Found empty disk %r, configuring it" % empty_disk
      prepareDisk(empty_disk)
    else:
      print "No empty disk found, trying to find the biggest possible free space"
      findAndPrepareFreeSpace()
  deactivateSwap()
  mount_point = '/mnt/' + LABEL
  if not os.path.exists(mount_point):
    os.mkdir(mount_point)
  os.chmod(mount_point, 0755)
  print 'Reconfiguring fstab'
  prepareSlapOSPartition(mount_point)
  print "Configuring fallback grub information"
  configureGrub(mount_point, slapos_label_file)
  configureNtp()

if __name__ == '__main__':
  try:
    run()
  except:
    import traceback
    sleep = 120
    print "There was uncaught issue with SlapOS configuration"
    print "System will be restarted in %ss" % sleep
    print "Below traceback might be useful:"
    traceback.print_exc()
    time.sleep(sleep)
    callWithIgnore(['init', '6'])
