#!/usr/bin/python
"""tsung python wrapper"""

import argparse
import logging
import sys
import tsung
import functools

def check_positive(value):
    ivalue = int(value)
    if ivalue <= 0:
         raise argparse.ArgumentTypeError("%s is an invalid positive int value" % value)
    return ivalue

def check_range(arg, minval, maxval):
	ivalue = int(arg)
	imin = int(minval)
	imax   = int(maxval)
	if ivalue < imin or ivalue > imax:
		raise argparse.ArgumentTypeError("%s is not in range %s - %s" % (ivalue, imin, imax))
	return ivalue

check_cpu = functools.partial(check_range, minval=1, maxval=30)
check_users = functools.partial(check_range, minval=1, maxval=70000)
check_maxusers = functools.partial(check_range, minval=1, maxval=100000)

logger = logging.getLogger()
logging.basicConfig(level=logging.INFO, stream=sys.stdout, format='%(levelname)s %(message)s')

def parse_arguments():
    parser = argparse.ArgumentParser()

    # basic arguments
    parser.add_argument("host", help="destination ip or hostname")
    parser.add_argument("--method", "-m", default="GET", choices=["GET", "POST"], help="HTTP method")
    parser.add_argument("--uri", default="/", help="destination uri")
    parser.add_argument("--port", "-p", default=80, help="destination port")
    parser.add_argument("--ssl", action="store_true", default=False, help="use secure connection")
    parser.add_argument("--clients", default="localhost", help="source ip or hosts comma-separated")
    parser.add_argument("--content", default="", help="request content; use file:// for load content from file")
    parser.add_argument("--cookies", default="", help="custom cookies in request")

    # performance arguments
    parser.add_argument("--users", "-u", default=1000, type=check_users, metavar='{1-70000}' ,help="concurent users")
    parser.add_argument("--requests", "-r", default=1, type=check_positive, metavar='{Positive number}', help='request per users')
    parser.add_argument("--maxusers", default=10000, type=check_maxusers, metavar='{500-100000}' ,help="maximum sockets per single process; but don't forget to raise the limit of the OS with ulimit -n")
    parser.add_argument("--cpu", default=4, type=check_cpu, metavar='{1-30}' ,help="maximum cpu to use")
    parser.add_argument("--controllervm", action="store_true", default=False, help="use controller vm; best performance with 'False' but fails during distributed task")

    # generic arguments
    parser.add_argument("--debug", action="store_true", default=False, help="debug output")
    parser.add_argument("--logtraffic", default="false", choices=["true", "false", "light", "protocol", "protocol_local"], help="tsung traffic dump")
    parser.add_argument("--dump", action="store_true", default=False, help="dump generator config")
    parser.add_argument("--dev", help="output iface; defaults to iface with default gw; eg. eth0")

    # timings
    parser.add_argument("--time", help="generate packets for specified time; eg. 2m3s")

    # template
    parser.add_argument("--template", help="custom template path")

    # logging
    parser.add_argument("--logdir", help="custom logdir, default ~/.tsung/log")

    return parser.parse_args()

def main():
	"""main"""
	args = parse_arguments()
	if args.debug:
		logger.setLevel(logging.DEBUG)

	runner = tsung.runner.Runner(args.__dict__)
	logger.debug("fields: %s", runner.fields)

	if args.dump:
		print runner.compile()
		ret = 0
	else:
		ret = runner.run()

	return ret

if __name__ == "__main__":
	sys.exit(main())
