#!@PYTHON@
#
# set ts=4, sw=4
#
# Author: Bas van der Vlies <basv@sara.nl>
# Date  : 16 February 2002
#
# Tester: Walter de Jong <walter@sara.nl>
#
# SVN info
#  $Id: pxeconfig.in 108 2007-10-03 10:40:30Z 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 

 [-r|--remove] -f|--file <filename>] [hostname(s)
 [-r|--remove] -i|--interactive
 [-r|--remove] -n|--net <C-net> -s|--start <number> -e|--end <number> -f|--file <filename>
 [-r|--remove] -b|--basename <string> s|--start <number> -e|--end <number> -f|--file <filename>
 -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
import ConfigParser

# DEBUG
#
DEBUG=0

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

SHORTOPT_LIST='be:f:hin:s:rV'
LONGOPT_LIST=['basename=', 'debug', 'end=', 'file=', 'help', 'interactive', 'net=', 'start=', 'remove', 'version' ]

def ReadConfig(file):
	"""
	Parse the config file
	"""
	if not os.path.isfile(file):
		print 'File %s does not exist' %file
		sys.exit(1)

	config = ConfigParser.RawConfigParser()
	try:
		config.read(file)
	except ConfigParser.MissingSectionHeaderError,detail:
		print detail
		sys.exit(1)

	# Not yet uses
	#
	#projects = {}
	#for section in config.sections():
	#	projects[section] = {}
	#	for option in  config.options(section):
	#	projects[section][option] = config.get(section, option) 

	stanza = config.defaults() 
	return stanza

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)

  # sort the files
  #
  files.sort()
 
  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, network, option_str=''):
	"""
	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 option_str, detail
		sys.exit(1)

	if not network:
		return n

	# Check if it is a correct network value
	#
	if 0 <= n <= 255:
		return n
	else:
		if option_str:
			number = option_str +' : ' + number 
			
		print '%s is not a valid network 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, True) 
	
	end = raw_input('Ending number: ') 
	end = check_number(end, True) 
	
	pxe_filename = select_pxe_configfile() 
	
	binfo[NETWORK] = naddr
	binfo[START] = start
	binfo[END] = end
	binfo[FILENAME] = pxe_filename 
	
	if DEBUG:
		print network, binfo

	manage_links(binfo)

def check_command_line(binfo, hostnames):
	"""
	Do you we have the right and propper values 
	"""
	### check_filename
	#
	try:
		if not os.path.isfile(os.path.join(PXE_CONF_DIR, binfo[FILENAME])):
			print '%s: Filename does not exists' %binfo[FILENAME]
			sys.exit(1)
	except KeyError, detail:
		binfo[FILENAME] = select_pxe_configfile() 

	if hostnames:
		host_2_net(hostnames, binfo)
		sys.exit(0)
		
	if binfo.has_key(BASENAME) and binfo.has_key(NETWORK):
		print __doc__
		print "The option -n/--net and -b/--basename are mutually exclusive"
		sys.exit(1)

	if binfo.has_key(BASENAME):
		network_number = False
		create_links =  base_2_net

	elif binfo.has_key(NETWORK):
		network_number = True
		create_links = manage_links
	else:
		print __doc__
		sys.exit(1)

	try:
		binfo[START] = check_number(binfo[START], network_number, '-s/--start')
	except KeyError, detail:
		print __doc__
		print '-s/--start is missing on the command line' %(detail)
		sys.exit(1)

	try:
		binfo[END] = check_number(binfo[END], network_number, '-e/--end')
	except KeyError, detail:
		print __doc__
		print '-e/--end is missing on the command line' %(detail)
		sys.exit(1)

	if DEBUG:
		print binfo

	create_links(binfo)


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.
	"""
	try:
		opts, args = getopt.gnu_getopt(argv[1:], SHORTOPT_LIST, LONGOPT_LIST)
	except getopt.error, detail:
		print __doc__
		print detail
		sys.exit(1) 
	
	global DEBUG

	# if nothing is specified then print usage and exit
	#
	if not opts and not args:
		print __doc__
		sys.exit(1)

	# Check given options
	#
	for opt,value in opts:
			
		if opt in ['-b', '--basename']:
			binfo[BASENAME] = value
			
		elif opt in ['--debug']:
			DEBUG = 1 
			
		elif opt in ['-e', '--end']: 
			binfo[END] = value 
			
		elif opt in ['-f', '--file']:
			binfo[FILENAME] = value 
			
		elif opt in ['-h', '--help']:
			print __doc__
			sys.exit(0) 

		elif opt in ['-i', '--interactive']:
			interactive(binfo)
			sys.exit(0) 

		elif opt in ['-n', '--net']:
			network = value
			binfo[NETWORK] = convert_network(value, opt) 
			
		elif opt in ['-r', '--remove']:
			binfo[REMOVE] = 1 
			
		elif opt in ['-s', '--start']:
			binfo[START] = value 

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

	check_command_line(binfo, args)

def host_2_net(hosts, binfo):
	"""
	Convert hostsname to a net address that can be handled by manage_links function
	"""
	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 base_2_net(binfo):
	"""
	Construct hostname(s) from the supplied basename and start and end numbers
	"""
	if binfo[START] >= binfo[END]:
		print __doc__
		print "Supplied wrong values for start (%d) and end (%d)" %(binfo[START], binfo[END])
		sys.exit(1)

	hostnames = list()
	for i in xrange(binfo[START], binfo[END] + 1):
		hostname = '%s%d' %(binfo[BASENAME], i)
		if DEBUG:
			print 'host = %s, Basename = %s, number = %d' %(hostname, binfo[BASENAME], i)
		hostnames.append(hostname)

	host_2_net(hostnames,binfo)
		
def main():
	# A dictionary holding the boot info
	#
	global DEBUG 
	
	bootinfo = {}
	bootinfo[REMOVE] = 0 
	
	configfile = '@pxeconfig_conf@'
	settings = ReadConfig(configfile) 
	
	try:
		PXE_CONF_DIR = settings['pxe_config_dir']
		if not DEBUG:
			DEBUG = int(settings['debug'])
	except KeyError:
		pass 

	check_args(sys.argv, bootinfo) 
	
if __name__ == '__main__':
	main()