#!/usr/bin/env python
#
# Author: Bas van der Vlies <basv@sara.nl>
# Date  : 16 February 2002
#
# Tester: Walter de Jong <walter@sara.nl>
#
# SVN info
#  $Id: pxeconfig 60 2005-10-31 11:50:24Z bas $
#
# Copyright (C) 2002
#
# This file is part of the pxeconfig utils
#
# 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 2, 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
#
"""
Usage: pxeconfig 

 [-d|--directory <pxe_config_dir>]
 [-f|--file <filename>] [hostname(s)]
 [-i|--interactive]
 [-n|--net <C-net> -s|--start <number> -e|--end <number> -f|--file <filename>]
 [-r|--remove]
 [-V|--version]

With this program you can configure which PXE configuration file a node
will use when it boots. The program can be started interactivly or via
command line options.

When this program is started interactive it will ask the following questions:
  1) Network address (Class C-network address only)
  2) Starting number
  3) Ending number
  4) Which PXE config file to use 

For example, if the answers are: 
  1) 10.168.44
  2) 2
  3) 4
  4) default.node_install

Then the result is that is create symbolic links in /tftpboot/pxelinux.cfg:
  0AA82C02 ----> default.node_install
  0AA82C03 ----> default.node_install
  0AA82C04 ----> default.node_install
"""

import string
import sys
import os
import glob
import getopt
import socket

# DEBUG
#
DEBUG=0

# Constants
#
PXE_CONF_DIR='/tftpboot/pxelinux.cfg'
NETWORK='network'
FILENAME='filename'
START='start'
END='end'
VERSION='0.6.2'
REMOVE='remove'
INTERACTIVE='interactive'

SHORTOPT_LIST='hVd:n:s:e:f:ri'
LONGOPT_LIST=['help', 'version', 'directory=', 'net=', 'start=', 'end=', 'file=', 'remove', 'debug' , 'interactive' ]

def select_pxe_configfile():
  """
  Let user choose which pxeconfig file to use.
  """

  os.chdir(PXE_CONF_DIR)

  # Try to determine to which file the default symlink points, and 
  # if it exists, remove it from the list.
  #
  try:
    default_file = os.readlink('default')
  except OSError:
    default_file = None
    pass

  files = glob.glob('default.*')
  if not files:
    print 'There are no pxe config files starting with: default.'
    sys.exit(1)

  if default_file:
    files.remove(default_file)
 

  print 'Which pxe config file must we use: ?' 
  i = 1   
  for file in files:
    print "%d : %s" %(i,file)
    i = i +1

  while 1:
    index = raw_input('Select a number: ')
    try:
      index = int(index)
    except ValueError:
      index = len(files) + 1

    # Is the user smart enough to select
    # the right value??
    #
    if 0 < index <= len(files): break

  return files[index-1]

def manage_links(dict):
  """
  Create the links in the PXE_CONF_DIR,
    list : A list containing: network hex address, pxe config file, 
           start number, end number
  """
  os.chdir(PXE_CONF_DIR)
  naddr = dict[NETWORK]
  pxe_filename = dict[FILENAME]
  for i in range(dict[START], dict[END]+1):
    haddr = '%s%02X' %(naddr, i)

    if dict[REMOVE]:
       if DEBUG:
          print 'removing %s/%s' %(PXE_CONF_DIR, haddr)
       if os.path.exists(haddr):
          os.unlink(haddr)
    else:
       if DEBUG:
          print 'linking %s to %s' %(haddr, pxe_filename)
       if os.path.exists(haddr):
          os.unlink(haddr)
       os.symlink(pxe_filename, haddr)

def convert_network(net, prefix=''):
  """
  This function checks if the give network is a Class C-network and will
  convert the network address to a hex address if true.
  """
  d = string.split(net, '.')

  if len(d) != 3:
    if prefix:
      net = prefix + ' : ' + net
      
    print '%s is not a valid  C-class network address!!!' %net
    sys.exit(1)


  # Display right message for interactive/commandline
  if prefix:
    prefix = prefix + ' ' + net
  else:
    prefix = net

  # Check if we have valid network values
  r = ''
  for i in d:
    i = check_number(i, prefix)
    r = '%s%02X' %(r,i)

  if DEBUG:
    print r

  return r

def check_number(number, prefix=''):
  """
  This functions checks if the input is between 0 < number < 255:
     number : is a string
     prefix ; a string that must be displayed if an error occurs
  """
  try:
    n = int(number)
  except ValueError, detail:
    print prefix,detail
    sys.exit(1)

  if 0 <= n <= 255:
    return n
  else:

    if prefix:
      number = prefix +' : ' + number

    print '%s is not a valid number, must be between 0 and 255' %number
    sys.exit(1)

def interactive(binfo):

  print __doc__

  network = raw_input('Give network address (xxx.xxx.xxx): ') 
  naddr = convert_network(network)

  start = raw_input('Starting number: ') 
  start = check_number(start)

  end = raw_input('Ending number: ') 
  end = check_number(end)

  pxe_filename = select_pxe_configfile()

  binfo[NETWORK] = naddr
  binfo[START] = int(start)
  binfo[END] = int(end)
  binfo[FILENAME] = pxe_filename

  if DEBUG:
    print network, binfo

def check_cmd_line(binfo):
  if len(binfo.keys()) != 5:
    print __doc__
    print 'Not enough arguments to create the links!!'
    sys.exit(1)

  # check_filename
  #
  if not os.path.isfile(os.path.join(PXE_CONF_DIR, binfo[FILENAME])):
    print '%s: Filename does not exists' %binfo[FILENAME]
    sys.exit(1)

def check_args(argv, binfo):
  """
  This function parses the command line options and returns the rest as
  an list of hostnames:
    argv     : a list of command line options.
    binfo    : returning a dict with the netinfo. if used non-interactively
    hostnames: the rest of the command lines options that are not-parseble.
  """
  global PXE_CONF_DIR
  global DEBUG

  try:
    opts, args = getopt.getopt(argv[1:], SHORTOPT_LIST, LONGOPT_LIST)
  except getopt.error, detail:
    print __doc__
    print detail
    sys.exit(1)

  # Check given options
  #
  for opt,value in opts:

    if opt in ['-d', '--directory']:
      if os.path.isdir(value):
        PXE_CONF_DIR = value
      else:
        print 'Directory %s does not exists\n' %value
	sys.exit(1)

    elif opt in ['-i', '--interactive']:
      binfo[INTERACTIVE] = 1

    elif opt in ['-n', '--net']:
      network = value
      binfo[NETWORK] = convert_network(value, opt)

    elif opt in ['-s', '--start']:
      binfo[START] = check_number(value, opt)

    elif opt in ['-e', '--end']:
      binfo[END] = check_number(value, opt)

    elif opt in ['-f', '--file']:
      binfo[FILENAME] = value

    elif opt in ['-V', '--version']:
      print VERSION
      sys.exit(0)

    elif opt in ['-r', '--remove']:
      binfo[REMOVE] = 1

    elif opt in ['--debug']:
      DEBUG = 1

    elif opt in ['-h', '--help']:
      print __doc__
      sys.exit(0)

  if args:
  	return args

def hosts_links(hosts, binfo):
	for host in hosts:
		try:
			addr = socket.gethostbyname(host)
		except socket.error,detail:
			print '%s not an valid hostname: %s' %(host,detail)
			sys.exit(1)
			
		net = string.splitfields(addr, '.')
		cnet = string.joinfields(net[0:3], '.')

		binfo[NETWORK] = convert_network(cnet)
		binfo[START] = int(net[3])
		binfo[END] =  int(net[3])
		manage_links(binfo)


def main():
  # A dictionary holding the boot info
  #
  bootinfo = {}
  bootinfo[REMOVE] = 0

  hostnames = check_args(sys.argv, bootinfo)

  if bootinfo[REMOVE]:
     bootinfo[FILENAME] = 'remove_boot_file'


  # If we supplied args then make links for the supplied hosts
  # else make links for class C-networks
  #
  if hostnames:
	if not bootinfo.has_key(FILENAME):
	      bootinfo[FILENAME] = select_pxe_configfile()
	hosts_links(hostnames,bootinfo)
  else:
  	if bootinfo.has_key(INTERACTIVE):    
		interactive(bootinfo)
  	else:
		check_cmd_line(bootinfo)
  	manage_links(bootinfo)

if __name__ == '__main__':
  main()