source: trunk/pxeconfig.in @ 125

#!@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 125 2008-04-13 16:46:59Z 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
  """
  str = 'manage_links'
  verbose(str)

  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 net_2_hex(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 = 'net_2_hex : %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 = net_2_hex(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_hex(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_host

        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] = net_2_hex(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_hex(hosts, binfo):
        """
        Convert hostname(s) to a net address that can be handled by manage_links function
        """
        str = 'host_2_hex: %s' %hosts
        verbose(str)

        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] = net_2_hex(cnet)
                binfo[START] = int(net[3])
                binfo[END] =  int(net[3])
                manage_links(binfo)

def base_2_host(binfo):
        """
        Construct hostname(s) from the supplied basename and start and end numbers
        """
        str = 'base_2_host'
        verbose(str)

        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_hex(hostnames,binfo)

def rack_2_net(binfo):
        """
        """
        str = 'rack_2_net'
        verbose(str)
        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_host(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)