#!/usr/bin/python

# Copyright (c)2010-2012 the Boeing Company.
# See the LICENSE file included in this distribution.
#
# author: Jeff Ahrenholz <jeffrey.m.ahrenholz@boeing.com>
#

"""
howmanynodes.py - This is a CORE script that creates network namespace nodes
having one virtual Ethernet interface connected to a bridge. It continues to
add nodes until an exception occurs. The number of nodes per bridge can be
specified.
"""

import datetime
import optparse
import shutil
import sys
import time

from core import constants
from core.misc import ipaddress
from core.netns import nodes
from core.session import Session

GBD = 1024.0 * 1024.0


def linuxversion():
    """ Return a string having the Linux kernel version.
    """
    f = open("/proc/version", "r")
    v = f.readline().split()
    version_str = " ".join(v[:3])
    f.close()
    return version_str


MEMKEYS = ("total", "free", "buff", "cached", "stotal", "sfree")


def memfree():
    """ Returns kilobytes memory [total, free, buff, cached, stotal, sfree].
        useful stats are:
            free memory = free + buff + cached
            swap used = stotal - sfree
    """
    f = open("/proc/meminfo", "r")
    lines = f.readlines()
    f.close()
    kbs = {}
    for k in MEMKEYS:
        kbs[k] = 0
    for l in lines:
        if l[:9] == "MemTotal:":
            kbs["total"] = int(l.split()[1])
        elif l[:8] == "MemFree:":
            kbs["free"] = int(l.split()[1])
        elif l[:8] == "Buffers:":
            kbs["buff"] = int(l.split()[1])
        elif l[:8] == "Cached:":
            kbs["cache"] = int(l.split()[1])
        elif l[:10] == "SwapTotal:":
            kbs["stotal"] = int(l.split()[1])
        elif l[:9] == "SwapFree:":
            kbs["sfree"] = int(l.split()[1])
            break
    return kbs


# node list (count from 1)
nodelist = [None]
switchlist = []


def main():
    usagestr = "usage: %prog [-h] [options] [args]"
    parser = optparse.OptionParser(usage=usagestr)
    parser.set_defaults(waittime=0.2, numnodes=0, bridges=0, retries=0,
                        logfile=None, services=None)

    parser.add_option("-w", "--waittime", dest="waittime", type=float,
                      help="number of seconds to wait between node creation" \
                           " (default = %s)" % parser.defaults["waittime"])
    parser.add_option("-n", "--numnodes", dest="numnodes", type=int,
                      help="number of nodes (default = unlimited)")
    parser.add_option("-b", "--bridges", dest="bridges", type=int,
                      help="number of nodes per bridge; 0 = one bridge " \
                           "(def. = %s)" % parser.defaults["bridges"])
    parser.add_option("-r", "--retry", dest="retries", type=int,
                      help="number of retries on error (default = %s)" % \
                           parser.defaults["retries"])
    parser.add_option("-l", "--log", dest="logfile", type=str,
                      help="log memory usage to this file (default = %s)" % \
                           parser.defaults["logfile"])
    parser.add_option("-s", "--services", dest="services", type=str,
                      help="pipe-delimited list of services added to each "
                           "node (default = %s)\n(Example: zebra|OSPFv2|OSPFv3|"
                           "IPForward)" % parser.defaults["services"])

    def usage(msg=None, err=0):
        sys.stdout.write("\n")
        if msg:
            sys.stdout.write(msg + "\n\n")
        parser.print_help()
        sys.exit(err)

    options, args = parser.parse_args()

    for a in args:
        sys.stderr.write("ignoring command line argument: %s\n" % a)

    start = datetime.datetime.now()
    prefix = ipaddress.Ipv4Prefix("10.83.0.0/16")

    print "Testing how many network namespace nodes this machine can create."
    print " - %s" % linuxversion()
    mem = memfree()
    print " - %.02f GB total memory (%.02f GB swap)" % (mem["total"] / GBD, mem["stotal"] / GBD)
    print " - using IPv4 network prefix %s" % prefix
    print " - using wait time of %s" % options.waittime
    print " - using %d nodes per bridge" % options.bridges
    print " - will retry %d times on failure" % options.retries
    print " - adding these services to each node: %s" % options.services
    print " "

    lfp = None
    if options.logfile is not None:
        # initialize a csv log file header
        lfp = open(options.logfile, "a")
        lfp.write("# log from howmanynodes.py %s\n" % time.ctime())
        lfp.write("# options = %s\n#\n" % options)
        lfp.write("# numnodes,%s\n" % ",".join(MEMKEYS))
        lfp.flush()

    session = Session(1)
    switch = session.add_object(cls=nodes.SwitchNode)
    switchlist.append(switch)
    print "Added bridge %s (%d)." % (switch.brname, len(switchlist))

    i = 0
    retry_count = options.retries
    while True:
        i += 1
        # optionally add a bridge (options.bridges nodes per bridge)
        try:
            if 0 < options.bridges <= switch.numnetif():
                switch = session.add_object(cls=nodes.SwitchNode)
                switchlist.append(switch)
                print "\nAdded bridge %s (%d) for node %d." % (switch.brname, len(switchlist), i)
        except Exception, e:
            print "At %d bridges (%d nodes) caught exception:\n%s\n" % (len(switchlist), i - 1, e)
            break

        # create a node
        try:
            n = session.add_object(cls=nodes.LxcNode, name="n%d" % i)
            n.newnetif(switch, ["%s/%s" % (prefix.addr(i), prefix.prefixlen)])
            n.cmd([constants.SYSCTL_BIN, "net.ipv4.icmp_echo_ignore_broadcasts=0"])
            if options.services is not None:
                session.services.add_services(n, "", options.services.split("|"))
                session.services.boot_services(n)
            nodelist.append(n)
            if i % 25 == 0:
                print "\n%s nodes created " % i,
                mem = memfree()
                free = mem["free"] + mem["buff"] + mem["cached"]
                swap = mem["stotal"] - mem["sfree"]
                print "(%.02f/%.02f GB free/swap)" % (free / GBD, swap / GBD),
                if lfp:
                    lfp.write("%d," % i)
                    lfp.write("%s\n" % ",".join(str(mem[x]) for x in MEMKEYS))
                    lfp.flush()
            else:
                sys.stdout.write(".")
            sys.stdout.flush()
            time.sleep(options.waittime)
        except Exception, e:
            print "At %d nodes caught exception:\n" % i, e
            if retry_count > 0:
                print "\nWill retry creating node %d." % i
                shutil.rmtree(n.nodedir, ignore_errors=True)
                retry_count -= 1
                i -= 1
                time.sleep(options.waittime)
                continue
            else:
                print "Stopping at %d nodes!" % i
                break

        if i == options.numnodes:
            print "Stopping at %d nodes due to numnodes option." % i
            break
        # node creation was successful at this point
        retry_count = options.retries

    if lfp:
        lfp.flush()
        lfp.close()

    print "elapsed time: %s" % (datetime.datetime.now() - start)
    print "Use the core-cleanup script to remove nodes and bridges."


if __name__ == "__main__":
    main()
