source: trunk/src/pxeconfig.py @ 202

# 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.py 202 2012-08-22 11:44:20Z 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 [-f,--filename <name>] <hosts|mac_addresses>

Specifying hosts or mac addresses:
   To specify a range use the [] to indicate a range,
   a couple of examples:

   The first five nodes of rack 16
    - gb-r16n[1-5]

   The first five nodes and node 12 and 18 of rack 16 to 20
      - gb-r[16-20]n[1-5,12,18]

   The first five nodes de in rack 16 with padding enabled 
    - gb-r[16]n[01-5]

   Host with mac address 00:19:b9:de:21:47 and first five node of rack 15
    - 00:19:b9:de:21:47 gb-r15n[1-5]
      
The ranges ([]) are not only limited to numbers, letters can also be used.

With this program you can configure which PXE configuration file a node
will use when it boots.

See following link for usage and examples:
 - https://subtrac.sara.nl/oss/pxeconfig/wiki/PxeUsage
"""

import string
import sys
import os
import glob
import getopt
import socket
#import ConfigParser
import re

# import from the sara_python_modules
import AdvancedParser
from pxe_global import *

# Constants
#
PXE_CONF_DIR = '/tftpboot/pxelinux.cfg'
NETWORK      = 'network'
START        = 'start'
END          = 'end'
VERSION      = '3.1.0'

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:
        error =  'There are no pxe config files starting with: default.'
        raise PxeConfig, error

    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(haddr, options):
    """
    Create the links in the PXE_CONF_DIR,
    list : A list containing: network hex address, pxe config file, 
           start number, end number
    """
    if options.VERBOSE:
        print 'manage_links(%s)' %(haddr)
    
    os.chdir(PXE_CONF_DIR)
    pxe_filename = options.filename
    
    if options.REMOVE:
        if options.DEBUG or options.DRY_RUN or options.VERBOSE:
            print 'removing %s/%s' %(PXE_CONF_DIR, haddr)

        if os.path.exists(haddr) and not options.DRY_RUN:
            os.unlink(haddr)

    else:
        if options.DEBUG or options.DRY_RUN or options.VERBOSE:
            print 'linking %s to %s' %(haddr, pxe_filename)

        if not options.DRY_RUN:
            if os.path.exists(haddr):
                os.unlink(haddr) 
            os.symlink(pxe_filename, haddr)

def net_2_hex(net, options):
    """
    This function checks if the give network is a Class C-network and will
    convert the network address to a hex address if true.
    """
    if options.DEBUG:
        str = 'net_2_hex : %s' %(net)
        print 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:
        r = '%s%02X' %(r, int(i))

    if options.DEBUG:
        print 'C-network converted to hex: ', r

    return r

def host_2_hex(host, options):
    """
    Convert hostname(s) to a net address that can be handled by manage_links function
    """
    if options.DEBUG:
        str = 'host_2_hex: %s' %hosts
        print str

    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], '.')

    if options.SCRIPT_HOOK:
        if options.DEBUG or options.DRY_RUN or options.VERBOSE:
            print 'Executing client script hook: %s with arg: %s' %(options.SCRIPT_HOOK, addr)
        if not options.DRY_RUN:
            cmd = '%s %s' %(options.SCRIPT_HOOK, addr)
            os.system(cmd)

    haddr = '%s%02X' %(net_2_hex(cnet, options), int(net[3])) 
    manage_links(haddr, options)


def mac_2_hex(mac_addr, options):
    """
    Convert mac address to pxeconfig address
    """
    if options.DEBUG:
        str = 'mac_2_hex: %s' %mac_addr
        print mac_addr

    haddr = '01-%s' %(mac_addr.replace(':', '-').lower())
    manage_links(haddr, options)

def add_options(p):
    """
    add the default options
    """
    p.set_defaults( 
        DEBUG   = False,
        VERBOSE = False,
        DRY_RUN = False,
        REMOVE  = False,
        VERSION  = False,
        SCRIPT_HOOK = False,
    )

    p.add_option('-d', '--debug',
        action = 'store_true',
        dest   = 'DEBUG',
        help   = 'Toggle debug mode (default : False)'
    )

    p.add_option('-f', '--filename',
        action = 'store',
        dest   = 'filename',
        help   = 'Specifies which PXE filename must be use'
    )

    p.add_option('-n', '--dry-run',
        action = 'store_true',
        dest   = 'DRY_RUN',
        help   = 'Do not execute any command (default : False)'
    )

    p.add_option('-r', '--remove',
        action = 'store_true',
        dest   = 'REMOVE',
        help   = 'Removes the PXE filename for a host(s)'
    )

    p.add_option('-v', '--verbose',
        action = 'store_true',
        dest   = 'VERBOSE',
        help   = 'Make the program more verbose (default: False)'
    )

    p.add_option('-V', '--version',
        action = 'store_true',
        dest   = 'VERSION',
        help   = 'Print the program version number and exit'
    )

def parser(argv, config, defaults): 
    """
    Make use of sara advance parser module. It can handle regex in hostnames
    """
    parser = AdvancedParser.AdvancedParser(usage=__doc__)
    add_options(parser)

    options, args = parser.parse_args() 

    if options.VERSION:
        print VERSION
        sys.exit(0) 

    if not args:
        parser.print_help()
        sys.exit(0)

    # debug can be specified by the command line or options file
    #
    if not options.DEBUG:
        try:
            if defaults['debug']:
                options.DEBUG = int(defaults['debug'])
        except KeyError:
            pass

    # Only check if we have specified an pxeconfig file if we did not
    # specify the REMOVE option
    #
    if not options.REMOVE:
        if options.filename:
            if not os.path.isfile(os.path.join(PXE_CONF_DIR, options.filename)):
                error =  '%s: Filename does not exists' %(options.filename)
                raise PxeConfig, error
        else:
            options.filename = select_pxe_configfile() 

    # ...
    try:
        options.SCRIPT_HOOK = defaults['client_script_hook']
    except KeyError, detail:
        pass

    if options.DEBUG:
        print args, options

    ##
    # Are the hosts wiht only mac addresses defined in the configuration file
    # or specified on the command line
    #
    mac_addr_re = re.compile('([a-fA-F0-9]{2}[:|\-]?){6}')

    for host in args:
        if host in config.sections():
            mac_addr = config.get(host, 'mac_address') 
            mac_2_hex(mac_addr, options)

        elif mac_addr_re.search(host):
            mac_2_hex(host, options)
        else:
            host_2_hex(host, options)

def main():
    # A dictionary holding the boot info
    #
    global PXE_CONF_DIR
    parser_config, default_settings = ReadConfig() 
    
    try:
        PXE_CONF_DIR = default_settings['pxe_config_dir']

    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

    parser(sys.argv, parser_config, default_settings)

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