#!@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 123 2008-04-11 15:10:23Z 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 

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

there are serveal ways to create a range of hostnames. they all start with:
 * -b,--base-name <string>

The different formats are:
  1  [opt] -H,--host-range <number>-<number> -f,--file <filename>
  2  [opt] -s,--start <number> -e,--end <number> -f,--file <filename>
  3  [opt] -R,--rack <string>-<string>  -N,--node <string>-<string> -f,--file <filename>

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

If leading zeros are used for <number>, then host names will be padded 
with zeros. --equal-width option will be enabled.

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
import re

# DEBUG
#
DEBUG=0
VERBOSE=0

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

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

# regex definition
#
# A regulare expression to parse string for rack/node
id_re = re.compile(r"""
	(?P<basename>[a-zA-Z-_]*)
	(?P<number>[0-9]+) 
     """, re.VERBOSE)

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

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
  """
  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] == True:
       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 'C-network in hex: ', 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_args(binfo, hostnames):
	"""
	Do you we have the right and propper values 
	"""
	### check_filename 
	#
	str = 'check_args: '
	verbose(str)	
	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):
		if binfo[RACK] and binfo[NODE]:
			create_links = rack_2_net
		else:
			set_padding(binfo)
			create_links =  base_2_net

	elif binfo.has_key(NETWORK):
		binfo[START] = check_number(binfo[START], True)
		binfo[END] = check_number(binfo[END], True)
		create_links = manage_links

	else:
		error = 'You have to specifiy -b,--basename or -n,--net'
		raise PxeConfig, error

	if DEBUG:
		print binfo

	create_links(binfo)

def set_padding(binfo):
	"""
	binfo          : boot info
	network_number : must we check if  start,end values are 
	                 valid network numbers
	return:
		- if equal_width is requested then the length will be set to end value
		- if start value length > 1 and start with a zero (0), width is set to
		  the end value
		- if end value starts with a zero (0), width will be set to the end 
		  value
	"""
	start_str = binfo[START]
	end_str = binfo[END]

	start = check_number(start_str, False)
	end = check_number(end_str, False)

	if binfo[EQUALWIDTH][0] == True:
		binfo[EQUALWIDTH][1] = len(end_str) 

	elif len(start_str) > 1 and start_str[0] == '0':
		binfo[EQUALWIDTH] = [ True, len(end_str) ]

	elif end_str[0] == '0':
		binfo[EQUALWIDTH] = [ True, len(end_str) ]

	binfo[START] = start
	binfo[END] = end


def parse_number_range(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

	binfo[START] = l[0]
	binfo[END] = l[1]

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

	l = arg.split('-')
	if len(l) < 2:
		error =  '%s : range syntax not valid,eg <string>-<string>)' %(arg)
		raise PxeConfig, error
	
	binfo[id] = dict()
	binfo[id][EQUALWIDTH] = [False, 0]

	i = 0
	for item in l: 
		result = id_re.match(item)

		if result:
			basename = result.group('basename')
			number = result.group('number')

			if DEBUG:
				print 'basename = %s, number = %s' %(basename, number)
	
			if i == 0:
				binfo[id][BASENAME] = basename
				binfo[id][START] = number
				i += 1
			else:
				binfo[id][END] = number

		else:
			error = '%s : string syntax is not valid, eg: <alpa>*<digit>+' %(item) 
			raise PxeConfig, error

        set_padding(binfo[id])	
	
def parse_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)

	# init vars
	#
	hostrange = node = rack = None

	# 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, 0]

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

		elif opt in ['-H', '--host-range']:
			hostrange = value

		elif opt in ['-N', '--node']:
			node = value

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

	if node and rack:
		 	parse_string_range(binfo, node, NODE)
		 	parse_string_range(binfo, rack, RACK)
	elif hostrange:
		 	parse_number_range(binfo, hostrange)
		

	check_args(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

	hostnames = list()
	for i in xrange(start, end + 1):
		if binfo[EQUALWIDTH][0] == True:
			hostname = '%s%0*d' %(binfo[BASENAME],  binfo[EQUALWIDTH][1], 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 rack_2_net(binfo):
	"""
	"""
	basename = binfo[BASENAME]
	start = binfo[RACK][START]
	end = binfo[RACK][END]


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

	for i in xrange(start, end + 1):
		if binfo[RACK][EQUALWIDTH][0] == True:
			new_base = '%s%s%0*d' %(basename,  binfo[RACK][BASENAME], binfo[RACK][EQUALWIDTH][1], i)
		else:
			new_base = '%s%s%d' %(basename, binfo[RACK][BASENAME], i)

		# Make copy and file in the appropiate values for creating/removing links
		#
		new_bootinfo = binfo[NODE].copy()
		new_bootinfo[BASENAME] = '%s%s' %(new_base, binfo[NODE][BASENAME])
		new_bootinfo[FILENAME] = binfo[FILENAME]
		new_bootinfo[REMOVE] = binfo[REMOVE]

		if DEBUG:
			print 'rack ', new_bootinfo

		base_2_net(new_bootinfo)

def main():
	# A dictionary holding the boot info
	#
	global DEBUG 
	global PXE_CONF_DIR
	
	bootinfo = {}
	bootinfo[NODE] = None
	bootinfo[RACK] = None
	bootinfo[REMOVE] = False 
	bootinfo[EQUALWIDTH] = [ False, 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 

	PXE_CONF_DIR = os.path.realpath(PXE_CONF_DIR)
	if not os.path.isdir(PXE_CONF_DIR):
		error =  'pxeconfig directory: %s does not exists' %(PXE_CONF_DIR)
		raise PxeConfig, error

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