#!/usr/bin/python

import sys
import os
import re
from optparse import OptionParser

from enum import Enum

from Tartarus.client import initialize
try:
    comm, argv = initialize()
except:
    e = sys.exc_info()[1]
    print ('Can\'t initialize runtime. '
           'Probably, this computer won\'t be joined to Tartarus domain.\n'
           'Error message: %s' % str(e))

from Tartarus.iface import DHCP, core
from Tartarus.system.ipaddr import IpAddr

class CommandLineError(Exception):
    def __init__(self, msg):
        super(CommandLineError, self).__init__(msg)

class CommandError(Exception):
    def __init__(self, msg, error_code=1):
        super(CommandError, self).__init__(msg)
        self.error_code = error_code

class Command(object):
    def run(self, argv):
        if len(argv) == 1 and argv[0] == 'help':
            self.help()
            return 0
        try:
            return self.do_cmd(argv)
        except CommandLineError, e:
            print '\033[91mError:\033[0m %s\n' % e
            self.help()
            return 1
        except CommandError, e:
            print '\033[91mError:\033[0m %s\n' % e
            return e.error_code
        except core.Error, e:
            print '\033[91mDHCP-server error:\033[0m %s\n' % e.reason
            return 1
    @classmethod
    def help(cls):
        name = cls.__name__.lower()
        print 'Usage: %s %s\n' % (sys.argv[0], name)
        print cls.__doc__
        print

class Start(Command):
    'start dhcpd daemon'
    def do_cmd(self, argv):
        prx = daemonPrx()
        prx.start()
        return 0

class Stop(Command):
    'stop dhcpd daemon'
    def do_cmd(self, argv):
        prx = daemonPrx()
        prx.stop()
        return 0

class Restart(Command):
    'stop and start dhcpd daemon'
    def do_cmd(self, argv):
        prx = daemonPrx()
        prx.stop()
        prx.start()
        return 0

class Status(Command):
    'show status of dhcpd daemon'
    def do_cmd(self, argv):
        prx = daemonPrx()
        if prx.running():
            print 'DHCP daemon is running'
        else:
            print 'DHCP daemon is not running'

class Subnet(Command):
    """manage IP subnets"""
    cidr_re = re.compile('^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}/\d{1,2}$')
    Action = Enum('LIST', 'ADD', 'REMOVE')
    def __init__(self):
        usage = (
                'Usage: %(prog)s subnet [-Ro]             - list mode\n'
                '       %(prog)s subnet -a <addr>/<mask>  - add mode\n'
                '       %(prog)s subnet -r <addr>/<mask>  - remove mode\n'
                )
        self.optp = OptionParser(usage=usage % { 'prog': sys.argv[0] })
        opt = self.optp.add_option
        Action = self.Action
        opt('-a', '--add',
                action='store_const', dest='action', const=Action.ADD,
                help='Create new range')
        opt('-r', '--remove',
                action='store_const', dest='action', const=Action.REMOVE,
                help='Remove range')
        opt('-R', '--show-ranges',
                action='store_true',
                help='Show subnet\'s ranges (only for list mode)')
        opt('-o', '--show-options',
                action='store_true',
                help='Show subnet options (only for list mode)')
        self.optp.set_defaults(action=Action.LIST)
    def help(self):
        self.optp.print_help()
    def do_cmd(self, args):
        Action = self.Action
        opts, args = self.optp.parse_args(args)
        {
                Action.LIST: self.do_list,
                Action.ADD: self.do_add,
                Action.REMOVE: self.do_remove,
        }[opts.action](opts, args)
    @classmethod
    def do_list(cls, opts, args):
        if len(args) != 0: raise CommandLineError('Wrong number of arguments')
        srv = serverPrx()
        for s in srv.subnets():
            cls.printSubnet(s, opts.show_options, opts.show_ranges)
    @classmethod
    def printSubnet(cls, subnet, options=False, ranges=False):
        cidr = subnet.cidr()
        print '%s' % (cidr)
        if options:
            for o in subnet.options().iteritems():
                print '  %s: %s' % o
        if ranges:
            for r in subnet.ranges():
                cls.printRange(r)
    @classmethod
    def printRange(cls, range):
        start, end = range.addrs()
        caps = range.caps()
        caps_str = ''
        if caps & DHCP.KNOWN: caps_str += 'k'
        if caps & DHCP.UNKNOWN: caps_str += 'u'
        print '  %s-%s CAPS: %s' % (start, end, caps_str)
    @classmethod
    def do_add(cls, opts, args):
        if len(args) != 1: raise CommandLineError('Wrong number of arguments')
        srv = serverPrx()
        srv.addSubnet(args[0])
    @classmethod
    def do_remove(cls, opts, args):
        if len(args) != 1: raise CommandLineError('Wrong number of arguments')
        cidr = args[0]
        if not cls.cidr_re.match(cidr):
            raise CommandLineError('Wrong subnet declaration: "%s"' % cidr)
        addr, mask = cidr.split('/')
        srv = serverPrx()
        subnet = srv.findSubnet(addr)
        if not subnet: raise CommandError('Subnet %s is not found on server' % cidr)
        srv.delSubnet(subnet.id())

class Range(Command):
    """manage IP ranges is subnets"""
    Action = Enum('LIST', 'ADD', 'REMOVE', 'MODIFY')
    def __init__(self):
        usage = (
                'Usage: %(prog)s range [<addr>]|[<addr>/<mask>] [-o]   - list mode\n'
                '       %(prog)s range -a <start> <end> <caps>         - add mode\n'
                '       %(prog)s range -r <addr>                       - remove mode\n'
                '       %(prog)s range -m <addr> <caps>                - modify mode'
                )
        self.optp = OptionParser(usage=usage % { 'prog': sys.argv[0] })
        opt = self.optp.add_option
        Action = self.Action
        opt('-a', '--add',
                action='store_const', dest='action', const=Action.ADD,
                help='Create new range')
        opt('-r', '--remove',
                action='store_const', dest='action', const=Action.REMOVE,
                help='Remove range')
        opt('-m', '--modify',
                action='store_const', dest='action', const=Action.MODIFY,
                help='Modify range capabilities')
        opt('-o', '--show-options',
                action='store_true',
                help='Show range options (only for list mode)')
        self.optp.set_defaults(action=Action.LIST)
    def help(self):
        self.optp.print_help()
    def do_cmd(self, args):
        Action = self.Action
        opts, args = self.optp.parse_args(args)
        {
                Action.LIST: self.do_list,
                Action.ADD: self.do_add,
                Action.REMOVE: self.do_remove,
                Action.MODIFY: self.do_modify
        }[opts.action](opts, args)
    @classmethod
    def do_list(cls, opts, args):
        srv = serverPrx()
        if len(args) == 0:
            for s in srv.subnets():
                ranges = [(IpAddr(r.addrs()[0]).int, r) for r in s.ranges()]
                ranges.sort(key=lambda r: r[0])
                for r in (r[1] for r in ranges):
                    cls.printRange(r, opts.show_options, s)
        elif len(args) == 1:
            if '/' in args[0]:
                s = srv.findSubnet(args[0].split('/')[0])
                if not s: return
                ranges = [(IpAddr(r.addrs()[0]).int, r) for r in s.ranges()]
                ranges.sort(key=lambda r: r[0])
                for r in (r[1] for r in ranges):
                    cls.printRange(r, opts.show_options, s)
            else:
                r = srv.findRange(args[0])
                if not r: return
                cls.printRange(r, opts.show_options)
    @classmethod
    def printRange(cls, range, options=False, subnet=None):
        start, end = range.addrs()
        if not subnet:
            srv = serverPrx()
            subnet = srv.findSubnet(start)
        caps = range.caps()
        caps_str = ''
        if caps & DHCP.KNOWN: caps_str += 'k'
        if caps & DHCP.UNKNOWN: caps_str += 'u'
        print '%s-%s (%s) CAPS: %s' % (start, end, subnet.cidr(), caps_str)
        if options:
            for o in range.options().iteritems():
                print '  %s: %s' % o
    @classmethod
    def do_add(cls, opts, args):
        if len(args) != 3: raise CommandLineError('Wrong number of arguments')
        start, end, caps = args
        c = 0
        if 'k' in caps: c |= DHCP.KNOWN
        if 'u' in caps: c |= DHCP.UNKNOWN
        srv = serverPrx()
        s = srv.findSubnet(start)
        if not s: raise CommandError('Can\'t find subnet for new range')
        s.addRange(start, end, c)
    @classmethod
    def do_remove(cls, opts, args):
        if len(args) != 1: raise CommandLineError('Wrong number of arguments')
        srv = serverPrx()
        r = srv.findRange(args[0])
        if not r: raise CommandError('Can\'t find range')
        s = srv.findSubnet(args[0])
        s.delRange(r.id())
    @classmethod
    def do_modify(cls, opts, args):
        if len(args) != 2: raise CommandLineError('Wrong number of arguments')
        addr, caps = args
        c = 0
        if 'k' in caps: c |= DHCP.KNOWN
        if 'u' in caps: c |= DHCP.UNKNOWN
        srv = serverPrx()
        r = srv.findRange(addr)
        if not r: raise CommandError('Can\'t find range')
        r.setCaps(c)

class Host(Command):
    """manage known hosts"""
    Action = Enum('LIST', 'ADD', 'REMOVE')
    def __init__(self):
        usage = (
                'Usage: %(prog)s host [<identity>]|[<hardware>] [-o]       - list mode\n'
                '       %(prog)s host -a <name> [<id>|ethernet <mac-addr>] - add mode\n'
                '       %(prog)s host -r <name>                            - remove mode\n'
                )
        self.optp = OptionParser(usage=usage % { 'prog': sys.argv[0] })
        opt = self.optp.add_option
        Action = self.Action
        opt('-a', '--add',
                action='store_const', dest='action', const=Action.ADD,
                help='Create new range')
        opt('-r', '--remove',
                action='store_const', dest='action', const=Action.REMOVE,
                help='Remove range')
        opt('-o', '--show-options',
                action='store_true',
                help='Show range options (only for list mode)')
        self.optp.set_defaults(action=Action.LIST)
    def help(self):
        self.optp.print_help()
    def do_cmd(self, args):
        Action = self.Action
        opts, args = self.optp.parse_args(args)
        {
                Action.LIST: self.do_list,
                Action.ADD: self.do_add,
                Action.REMOVE: self.do_remove,
        }[opts.action](opts, args)
    @classmethod
    def do_list(cls, opts, args):
        srv = serverPrx()
        if len(args) != 0: raise CommandLineError('Wrong number of arguments')
        for h in srv.hosts():
            cls.printHost(h, opts.show_options)
    @classmethod
    def printHost(cls, host, options=False):
        name = host.name()
        id = host.id()
        if id.type == DHCP.HostIdType.IDENTITY:
            str_id = 'id: %s' % id.value
        else:
            str_id = 'hw: %s' % id.value
        print '%s(%s)' % (name, str_id)
        if not options: return
        for i in host.options().iteritems():
            print '  %s: %s' % i
    @classmethod
    def do_add(cls, opts, args):
        if len(args) == 1:
            id = DHCP.HostId(DHCP.HostIdType.IDENTITY, args[0])
        elif len(args) == 2:
            id = DHCP.HostId(DHCP.HostIdType.IDENTITY, args[1])
        elif len(args) == 3 and args[1] == 'ethernet':
            id = DHCP.HostId(DHCP.HostIdType.HARDWARE, 'ethernet '+args[2])
        else: raise RuntimeError('Wrong arguments')
        srv = serverPrx()
        srv.addHost(args[0], id)
    @classmethod
    def do_remove(cls, opts, args):
        if len(args) != 1: raise CommandLineError('Wrong number of arguments')
        srv = serverPrx()
        srv.delHost(args[0])

class Options(Command):
    """manipulate DHCP options"""
    Action = Enum('LIST', 'SET', 'UNSET')
    subnet_re = Subnet.cidr_re
    range_re = re.compile('^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$')
    def __init__(self):
        usage = (
                'Usage: %(prog)s options [<scope>]                - list mode\n'
                '       %(prog)s options [<scope>] <key> <value>  - set mode\n'
                '       %(prog)s options -r [<scope>] <key>       - unset mode\n\n'
                '<scope> is one of:\n'
                '  - global for global options\n'
                '  - <ip-addr>/<mask> for subnet options\n'
                '  - <ip-addr> for ip-range options\n'
                '  - <hostname> for host options\n'
                )
        self.optp = OptionParser(usage=usage % { 'prog': sys.argv[0] })
        opt = self.optp.add_option
        Action = self.Action
        opt('-r', '--remove',
                action='store_const', dest='action', const=Action.UNSET,
                help='Unset option')
    def help(self):
        self.optp.print_help()
    def do_cmd(self, args):
        Action = self.Action
        opts, args = self.optp.parse_args(args)
        if opts.action is None:
            if len(args) > 1: opts.action = self.Action.SET
            else: opts.action = self.Action.LIST
        {
                Action.LIST: self.do_list,
                Action.SET: self.do_set,
                Action.UNSET: self.do_unset,
        }[opts.action](opts, args)
    @classmethod
    def do_list(cls, opts, args):
        if len(args) == 0: scope = 'global'
        elif len(args) == 1: scope = args[0]
        else: raise CommandLineError('Wrong number of arguments')
        opts = cls.proxyForScope(scope).options()
        cls.printOpts(opts)
    @classmethod
    def printOpts(cls, options):
        for o in options.iteritems():
            print '%s: %s' % o
    @classmethod
    def do_set(cls, opts, args):
        if len(args) == 2:
            scope = 'global'
            key, value = args
        elif len(args) == 3:
            scope, key, value = args
        else: raise CommandLineError('Wrong number of arguments')
        prx = cls.proxyForScope(scope)
        prx.setOption(key, value)
    @classmethod
    def do_unset(cls, opts, args):
        if len(args) == 1:
            scope = 'global'
            key = args[0]
        elif len(args) == 2:
            scope, key = args
        else: raise CommandLineError('Wrong number of arguments')
        prx = cls.proxyForScope(scope)
        prx.unsetOption(key)
    @classmethod
    def proxyForScope(cls, scope):
        srv = serverPrx()
        if scope == 'global':
            return srv
        if cls.subnet_re.match(scope):
            addr, _ = scope.split('/')
            subnet = srv.findSubnet(addr)
            return subnet
        if cls.range_re.match(scope):
            range = srv.findRange(scope)
            return range
        host = srv.getHost(scope)
        return host

class Main:
    def __init__(self, argv):
        self.__exec_name = sys.argv[0]
        self.__argv = argv
        self.__cmds = {}
        for cmd, cls in ((cmd.lower(), cls) for cmd, cls in globals().iteritems() if
                isinstance(cls, type)
                and issubclass(cls, Command)
                and cls is not Command
                ):
            self.__cmds[cmd] = cls
    def usage(self):
        print 'Usage: %s <command> <args>\n' % sys.argv[0]
        self.helpCmds()
    def helpCmds(self):
        print 'Valid commands are:'
        for cmd, doc in ((key, cmd.__doc__) for key, cmd in self.__cmds.iteritems()):
            print '  \033[92m%s\033[0m - %s' % (cmd, doc)
        print '\nTry run %s <command> --help for more information\n' % sys.argv[0]
    def __call__(self):
        if len(self.__argv) < 2:
            self.usage()
            return 1
        cmd = self.__argv[1]
        if cmd == '--help':
            self.usage()
            return 0
        cmd = self.__cmds.get(cmd, None)
        if cmd is None:
            print 'Error: %s - unknown command' % self.__argv[1]
            self.helpCmds()
            return 1
        argv = self.__argv[2:]
        cmd().run(argv)
        comm.destroy()
        return 0

def daemonPrx():
    prx = comm.stringToProxy('DHCP/Daemon')
    prx = DHCP.DaemonPrx.checkedCast(prx)
    if not prx: raise CommandError('Can\'t connect to server')
    return prx

def serverPrx():
    prx = comm.stringToProxy('DHCP/Server')
    prx = DHCP.ServerPrx.checkedCast(prx)
    if not prx: raise CommandError('Can\'t connect to server')
    return DHCP.ServerPrx.checkedCast(prx)

sys.exit(Main(argv)())

