source: trunk/pxeconfig.in @ 117

#!@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 117 2008-04-10 16:41:48Z 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
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:rwvV'
LONGOPT_LIST=[ 'basename=', 'debug', 'end=', 'equal-width', 
    'file=', 'help', 'interactive', 'net=', 'start=', 
        'remove', 'verbose', 'version', 'equal-width'
        ]

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

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:
                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):
                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
        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']:
                        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 ['-w', '--equal-width']:
                        binfo[EQUALWIDTH] = True 

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

                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:
                        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)

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

        hostnames = list()
        for i in xrange(binfo[START], binfo[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 

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

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