#!/usr/bin/python
#
# Version 1.1
#
# Scan the system for disk drives, and set up partitioning according
#   to drives found, and free space/Linux partitions on them
#
# Scan the kickstart arguments and use a package set accordingly.
#   Set interactive the same way.
#

import os
import sys
import parted

# for testing
if (os.path.exists('isys')):
    sys.path.append('isys')

sys.path.append('/usr/lib/anaconda')

import isys

# Set up the package set variables
small = """
@ Core
-lilo
-redhat-config-mouse
-redhat-config-network-tui
-anacron
-apmd
-aspell
-autofs
-devlabel
-dump
-fbset
-finger
-ftp
-irda-utils
-kernel-pcmcia-cs
-logwatch
-man-pages
-mtr
-nss_ldap
-pam_krb5
-pam_smb
-pax
-pidentd
-reiserfs-utils
-rp-pppoe
-rsh
-ext3utils
-jwhois
-setuptool
-sendmail
-specspo
-sudo
-stunnel
-up2date
-tcpdump
-lrzsz
-wvdial
-isdn4k-utils
-tk
m4
cpp
glibc-devel
libxml2
glibc-kernheaders
patch
dosfstools
libstdc++-devel
mtools
dialog
autoconf
automake
binutils
gcc
gcc-c++
libtool
make
rpm-build
vim-enhanced
elinks
openssh-server
openssh-clients
compat-libstdc++
libxml2-python
jfsutils
kernel-utils"""

rm = """
@ base-x
@ gnome-desktop
@ editors
@ graphical-internet
@ text-internet
@ server-cfg
@ web-server
@ mail-server
@ smb-server
@ mysql
@ sql-server
@ legacy-network-server
@ network-server
@ development-tools
@ admin-tools
@ system-tools
@ printing
sysstat"""

rm_as = """
@ ftp-server
@ dns-server
@ news-server"""

rm_64 = """
@ compat-arch-support
@ compat-arch-development"""

ws = "@ Everything"

inhouse = """
@ base-x
@ gnome-desktop
@ editors
@ graphical-internet
@ text-internet
@ sound-and-video
@ office
@ development-tools
@ system-tools
@ printing
rdesktop
kdebase
kdepim
gimp
xpdf"""

# Set up xconfig
basexconfig = 'xconfig --depth=16 --resolution=1024x768 --defaultdesktop=GNOME --monitor="Generic High Frequency SVGA, 1024x768 @ 70 Hz'
xonboot = ' --startxonboot'

# Set up other aliases
drives = isys.hardDriveDict()
driveList = drives.keys()
driveList.sort()
firstdisk = driveList[0]
otherdisks = driveList[1:]
drivesizes = {}
arch = (os.uname())[4]
data = 0
zantaz = 0

# Open files
packagefile = open('/tmp/packages', 'a')
partfile = open('/tmp/parts', 'a')
interactivefile = open('/tmp/interactive', 'a')
xconfigfile = open('/tmp/xconfig', 'a')
cmdline = open("/proc/cmdline", "r").read()

# Get Disk Sizes
for drive in driveList:
  # try to open and get size
  skip = 0
  deviceFile =  "/tmp/%s" % (drive,)
  isys.makeDevInode(drive, deviceFile)
  dev = parted.PedDevice.get(deviceFile)
  sizeMB = (float(dev.heads * dev.cylinders * dev.sectors) / (1024 * 1024) * dev.sector_size)
  drivesizes[drive] = sizeMB
  if drivesizes[drive] > 2097152.0:
    print "Found" ,drive ,"which is larger than 2TB.  Making a gpt partition table."
    disk = dev.disk_new_fresh(parted.disk_type_get('gpt'))
    disk.commit()
# Here we try to see if we can open the disk.  If we can't, make a new label on the disk.
  try:
    disk = parted.PedDisk.new(dev)
  except parted.error:
    print "Found disk" ,drive ,"with no partition table.  Making an msdos partition table."
    disk = dev.disk_new_fresh(parted.disk_type_get('msdos'))
    disk.commit()
                                                                                                                                                                                     
  os.remove(deviceFile)

# Check to see if this is Zantaz
if cmdline.find('zantaz') > 0:
  zantaz = 1
  print 'I found zantaz'

# Set up our basic partitions
partfile.write('part /boot --fstype ext3 --size 100 --asprimary --ondisk=' + firstdisk + '\n')
partfile.write('part swap --size 2048 --asprimary --ondisk=' + firstdisk + '\n')
# Is this Zantaz?
#if zantaz == 1:
#  partfile.write('part swap --size 2048 --ondisk=' + firstdisk + '\n')
#  partfile.write('part / --fstype ext3 --size 1024 --ondisk=' + firstdisk + '\n')
#  partfile.write('part /usr --fstype ext3 --size 4096 --ondisk=' + firstdisk + '\n')
#  partfile.write('part /var --fstype ext3 --size 4096 --ondisk=' + firstdisk + '\n')
#  partfile.write('part /opt --fstype ext3 --size 10240 --grow --ondisk=' + firstdisk + '\n')
#else:
# If our first disk is rather large (greater than 7GB), do things slightly different
#  if drivesizes[firstdisk] > 716800.0:
#    partfile.write('part / --fstype ext3 --asprimary --size 20480 --ondisk=' + firstdisk + '\n')
#    partfile.write('part /data --fstype xfs --asprimary --size 1 --grow --ondisk=' + firstdisk + '\n')
#    data = 1
#  else:
#    if cmdline.find('myth') > 0:
#      partfile.write('part / --fstype xfs --asprimary --size 10240 --grow --ondisk=' + firstdisk + '\n')
#    else:
#      partfile.write('part / --fstype ext3 --asprimary --size 10240 --grow --ondisk=' + firstdisk + '\n')
partfile.write('part / --fstype ext3 --asprimary --size 10240 --grow --ondisk=' + firstdisk + '\n')

# If small, don't do data partitions
if cmdline.find('small') == -1:
  for i in range(len(otherdisks)):
    disk = parted.PedDisk.new(parted.PedDevice.get("/dev/" + otherdisks[i]))
    p = disk.next_partition()
# Check to see if there is any free space or a Linux partition to use
    while p:
      if p.type & parted.PARTITION_METADATA:
        p = disk.next_partition(p)
        continue
# Lets make sure that if we have a free space partition, that it is at least 2 gigs before we use it
      if p.type & parted.PARTITION_FREESPACE and (p.geom.length * p.geom.dev.sector_size / 1024.0 / 1024.0) > 2048:
        if zantaz == 1:
          partfile.write('part /opt/data --fstype ext3 --asprimary --size 1 --grow --ondisk=' + otherdisks[i] + '\n')
          p = None
          continue
        else:
          if data == 0:
            partfile.write('part /data --fstype xfs --asprimary --size 1 --grow --ondisk=' + otherdisks[i] + '\n')
            p = None
            data = data + 1
            continue
          else:
            partfile.write('part /data' + str(data) + ' --fstype xfs --asprimary --size 1 --grow --ondisk=' + otherdisks[i] + '\n')
            p = None
            data = data + 1
            continue
            linpart = disk.next_partition()
            partrequested = 0
# So if we didn't find any free space, then check for Linux partitions
      try:
        if p.fs_type.name == 'ext3' or p.fs_type.name == 'ext2' or p.fs_type.name == 'jfs' or p.fs_type.name == 'xfs' or p.fs_type.name == 'reiserfs' or p.fs_type.name == 'linux-swap':
          if zantaz == 1:
            partfile.write('part /opt/data --fstype ext3 --asprimary --size 1 --grow --ondisk=' + otherdisks[i] + '\n')
            p = None
            continue
          else:
            if data == 0:
              partfile.write('part /data --fstype xfs --asprimary --size 1 --grow --ondisk=' + otherdisks[i] + '\n')
              data = data + 1
              p = None
              continue
            else:
              partfile.write('part /data' + str(data) + ' --fstype xfs --asprimary --size 1 --grow --ondisk=' + otherdisks[i] + '\n') 
              data = data + 1
              p = None
              continue
# Except AttributeError because the file system type is None?  Pass on this and continue on.
      except AttributeError:
        p = disk.next_partition(p)
        continue
      p = disk.next_partition(p) 
          

# Now check the cmdline for which package set and xconfig to use
if cmdline.find(' small ') > 0:
  packagefile.write(small)
  xconfigfile.write(basexconfig)
elif cmdline.find(' zantaz ') > 0:
  packagefile.write(small)
  xconfigfile.write(basexconfig + xonboot)
elif cmdline.find(' ws ') > 0:
  packagefile.write(ws)
  xconfigfile.write(basexconfig + xonboot)
elif cmdline.find(' inhouse ') > 0:
  packagefile.write(inhouse)
  xconfigfile.write(basexconfig + xonboot)
elif cmdline.find(' rm ') > 0:
  packagefile.write(rm)
  if cmdline.find('rhel4as') > 0 or cmdline.find('rhel4es') > 0:
    packagefile.write(rm_as)
  if cmdline.find('64') > 0:
    packagefile.write(rm_64)
  xconfigfile.write(basexconfig)
else:
  packagefile.write(rm)
  xconfigfile.write(basexconfig)

# Now check the cmdline for if we should be interactive or not.
if cmdline.find('int') > 0:
  interactivefile.write('interactive\n')

partfile.close()
packagefile.close()
interactivefile.close()
xconfigfile.close()

print 'Our partition table is as follows: \n'
print open("/tmp/parts", "r").read()
partfile.close()

################
# Changelog
#
#  Mon Feb 28 2005 Jesse Keating <jesse.k@pogolinux.com> 1.4
#  - Ported to RHEL4, fixed some search issues that would cause rhel3ws to trigger the WS package set.
#
#  Tue Oct 26 2004 Jesse Keating <jesse.k@pogolinux.com> 1.3
#  - Ported to RHEL3
#
#  Thu Sep 23 2004 Jesse Keating <jesse.k@pogolinux.com> 1.2
#  - Fixed >2TB issues
#
#  Wed Aug 25 2004 Jesse Keating <jesse.k@pogolinux.com> 1.1
#  - Added Zantaz partition scheme and option
#
#  Thu Jul 29 2004 Jesse Keating <jesse.k@pogolinux.com> 1.0
#  - Released to production