#!@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 119 2008-04-10 21:52:45Z 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 

 [optional] -f|--file <filename> <hostname(s)>
 [optional] -i|--interactive
 [optional] -n|--net <C-net> -s|--start <number> -e|--end <number> -f|--file <filename>
 [optional] -b|--basename <string> s|--start <number> -e|--end <number> -f|--file <filename>
 -V|--version

options:
 -r,--remove 
 -w,--equal-width

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
VERBOSE=0

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

SHORTOPT_LIST='b:e:f:hin:s:rwvH:V'
LONGOPT_LIST=[ 'basename=', 'debug', 'end=', 'equal-width', 
    'file=', 'help', 'host-range=', 'interactive', 'net=', 
    'start=', 'remove', 'verbose', 'version', 'equal-width'
	]

def verbose(str):
	if VERBOSE:
		print '%s'

class PxeConfig(Exception):
	def __init__(self, msg=''):
		self.msg = msg
		Exception.__init__(self, msg)

	def __repr__(self):
		return self.msg


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
  """
  print PXE_CONF_DIR
  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):
	"""
	This function checks if the give network is a Class C-network and will
	convert the network address to a hex address if true.
	"""
	str = 'convert_network : %s' %(net)
	verbose(str) 
	
	d = string.split(net, '.') 
	if len(d) != 3:
		error = '%s is not a valid  C-class network address' %(net)
		raise PxeConfig, error

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

	if DEBUG:
		print r

	return r

def check_number(number_str, network):
	"""
	number : a string. If string starts with a zero (0) then 
	         EQUALWIDTH wil be set.
	network: if true then number must ben between 0 < number < 255
	         else it must be a valid number. 
	"""

	try:
		n = int(number_str)
	except ValueError, detail:
		error = "%s : is not a valid number" %number_str
		raise PxeConfig, error

	if not network:
		return n

	# Check if it is a correct network value
	#
	if 0 <= n <= 255:
		return n
	else:
		error = '%s is not a valid network number, must be between 0 and 255' %n
		raise PxeConfig, error

def interactive(binfo):
	print __doc__ 
	

	while 1:
		network = raw_input('Give network address (xxx.xxx.xxx): ') 

		try:
			naddr = convert_network(network) 
			break
		except PxeConfig, detail:
			print '%s : not a valid C-class network number' %(network)
			continue

	while 1:
		start = raw_input('Starting number: ') 

		try:
			start = check_number(start, True) 
			break
		except PxeConfig, detail:
			print detail
			continue
		
	while 1:
		end = raw_input('Ending number: ') 

		try:
			end = check_number(end, True) 
			break
		except PxeConfig, detail:
			print detail
			continue
	
	
	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])):
			error =  '%s: Filename does not exists' %binfo[FILENAME]
			raise Pxeconfig, detail
	except KeyError, detail:
		if binfo[REMOVE] :
			binfo[FILENAME] = 'Does not matter'
		else:
			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):
		error =  "The option -n/--net and -b/--basename are mutually exclusive"
	 	raise PxeConfig, error	

	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:
		error = 'You have to specifiy -b,--basename or -n,--net'
		raise PxeConfig, error

	binfo[START] = check_number(binfo[START], network_number)
	binfo[END] = check_number(binfo[END], network_number)

	if DEBUG:
		print binfo

	create_links(binfo)

def set_padding(binfo, start, end):
	"""
	find out th
	"""
	if len(start) > 1 and start[0] == '0':
		binfo[EQUALWIDTH] = True
	elif end[0] == '0':
		binfo[EQUALWIDTH] = True

	if len(start) == len(end):
		a =1 	
		

def parse_hostrange(binfo, arg):
	"""
	Parse if arg is of format <digit-digit>, if it starts 
	with a zero (0) then set EQUALWIDTH
	"""
	str = 'parse_hostrange %s' %(arg)
	verbose(str)

	l = arg.split('-')
	if len(l) < 2:
		error =  'hostrange syntax not valid: %s (number-number)' %(arg)
		raise PxeConfig, error

	start = l[0]
	end = l[1]


	binfo[START] = check_number(start, False)
	binfo[END] = check_number(end, False)
	
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
	global VERBOSE

	# 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']:
			binfo[NETWORK] = convert_network(value) 
			
		elif opt in ['-r', '--remove']:
			binfo[REMOVE] = 1 
			
		elif opt in ['-s', '--start']:
			binfo[START] = value 

		elif opt in ['-w', '--equal-width']:
			binfo[EQUALWIDTH] = True 

		elif opt in ['-v', '--verbose']:
			VERBOSE = 1

		elif opt in ['-H', '--host-range']:
		 	parse_hostrange(binfo, value)
			
		elif opt in ['-V', '--version']:
			print VERSION
			sys.exit(0) 

	check_command_line(binfo, args)

def host_2_net(hosts, binfo):
	"""
	Convert hostname(s) 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:
			error =  '%s not an valid hostname: %s' %(host,detail)
			raise PxeConfig, error
			
		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
	"""
	start = binfo[START]
	end = binfo[END]

	if start > end:
		error = '%d >= %d : start value is greater then end value' %(start, end)
		raise PxeConfig, error

	if binfo[EQUALWIDTH]:
		width = len(str(end))

	hostnames = list()
	for i in xrange(start, end + 1):
		if binfo[EQUALWIDTH]:
			hostname = '%s%0*d' %(binfo[BASENAME], width, i)
		else:
			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 
	global PXE_CONF_DIR
	
	bootinfo = {}
	bootinfo[REMOVE] = False 
	bootinfo[EQUALWIDTH] = False 
	
	configfile = '@pxeconfig_conf@'
	settings = ReadConfig(configfile) 
	
	try:
		PXE_CONF_DIR = settings['pxe_config_dir']
		if not DEBUG:
			DEBUG = int(settings['debug'])

	except KeyError:
		pass 

	PXE_CONF_DIR = os.path.realpath(PXE_CONF_DIR)
	if not os.path.isdir(PXE_CONF_DIR):
		print 'pxeconfig directory: %s does not exists' %(PXE_CONF_DIR)
		sys.exit(1)

	try:
		check_args(sys.argv, bootinfo) 
	except PxeConfig, detail:
		print detail
		sys.exit(1)
	
if __name__ == '__main__':
	main()