#!/usr/bin/python
'''
Code licensed under the BSD License:
Author : emptyhua@gmail.com
Site : http://bluehua.org

'''

import sys
import threading
import signal
import time
import Queue
import pycurl 
import string
import re
import random

TOTAL_REQUESTS = 0
REQUEST_SUCCESS = 0
TOTAL_REQUEST_TIME = 0
TIMEOUT_COUNT = 0
REFUSED_COUNT = 0
NOT2XX_COUNT = 0

TIMEOUT = 3

CLIENT_NUM = 0
TOTAL_TIME = 0

GETRANGE = re.compile(r'(\{([\d]+)[\-,]([\d]+)(#random)?\})')

def void_write(s):
    pass

class Request(threading.Thread):

    def __init__(self, queue):
        threading.Thread.__init__(self)
        self.queue = queue


    def run(self):
        global TOTAL_REQUESTS,TOTAL_REQUEST_TIME,REQUEST_SUCCESS,TIMEOUT,TIMEOUT_COUNT,REFUSED_COUNT,NOT2XX_COUNT
        while True:
            url = self.queue.get()

            TOTAL_REQUESTS += 1
            start_time = time.clock()
            try:
                c = pycurl.Curl()
                c.setopt(pycurl.URL, url)
                c.setopt(pycurl.NOSIGNAL, 1)
                c.setopt(pycurl.WRITEFUNCTION, void_write)
                c.setopt(pycurl.TIMEOUT, TIMEOUT)
                c.perform()
                if c.getinfo(pycurl.HTTP_CODE) / 100 == 2:
                    REQUEST_SUCCESS += 1
                else:
                    NOT2XX_COUNT += 1
            except pycurl.error,ex:
                if ex[0] == 7:
                    REFUSED_COUNT += 1
                elif ex[0] == 28:
                    TIMEOUT_COUNT += 1



            TOTAL_REQUEST_TIME += time.clock()

            self.queue.task_done()

class Urlparser:
    numrange = None
    current_index = 0
    random= False

    def __init__(self, url):
        self.url = url
        regrt =  GETRANGE.findall(url)
        if len(regrt) > 0:
            regrt = regrt[0]
            self.url = self.url.replace(regrt[0], '~#range#~')
            self.numrange = (string.atoi(regrt[1]), string.atoi(regrt[2]))
            self.current_index = self.numrange[0]
            if regrt[3] == '#random':
                self.random = True

    def get(self):
        if self.numrange is None:
            return self.url
        else:
            rt = self.url.replace('~#range#~', str(self.current_index))
            if self.random:
                self.current_index = random.randint(self.numrange[0], self.numrange[1])
            else:
                if self.current_index == self.numrange[1]:
                    self.current_index = self.numrange[0]
                else:
                    self.current_index += 1 
            return rt

def do_request(url, client_num, req_time, timeout):
    global TOTAL_REQUESTS,REQUEST_SUCCESS,TOTAL_REQUEST_TIME,CLIENT_NUM,TOTAL_TIME,TIMEOUT,TIMEOUT_COUNT,REFUSED_COUNT,NOT2XX_COUNT
    TOTAL_REQUESTS = 0
    REQUEST_SUCCESS = 0
    TOTAL_REQUEST_TIME = 0
    CLIENT_NUM = client_num
    TOTAL_TIME = 0
    TIMEOUT_COUNT = 0
    REFUSED_COUNT = 0
    NOT2XX_COUNT = 0
    TIMEOUT = timeout
    
    urls = Urlparser(url)
    queue = Queue.Queue()

    for i in xrange(client_num):
        r = Request(queue)
        r.setDaemon(True)
        r.start()

    start_time = time.clock()

    while TOTAL_TIME <= req_time:
        if queue.empty():
            queue.put(urls.get())
        TOTAL_TIME = time.clock()

    queue.join()
    report()

def report():
    global TOTAL_REQUESTS,REQUEST_SUCCESS,TOTAL_REQUEST_TIME,CLIENT_NUM,TOTAL_TIME
    print
    print '%d clients, running %d sec, timeout setting %d sec.' % (CLIENT_NUM, TOTAL_TIME, TIMEOUT)
    print 'Request speed: %d/sec.' % (TOTAL_REQUESTS / int(TOTAL_TIME))
    print 'Requests: %d total ,%d susceed, average %d%%.' % (TOTAL_REQUESTS, REQUEST_SUCCESS, int(float(REQUEST_SUCCESS) / TOTAL_REQUESTS * 100))
    print '%d fail:' % (TOTAL_REQUESTS - REQUEST_SUCCESS)
    print 'timeout:%d' % TIMEOUT_COUNT
    print 'connect refused:%d' % REFUSED_COUNT
    print 'not 2xx status:%d' % NOT2XX_COUNT

def safe_exit(num, frame):
    report()
    exit()

if __name__ == '__main__':
    from optparse import OptionParser
    optparser = OptionParser( 'usage: %prog -c num -t num url' )
    optparser.add_option( '-c' , '--client' , dest='client_num' , type='int', default='1', help='client number. default:1' )
    optparser.add_option( '-t' , '--time' , dest='req_time' , type='int', default='30', help='total run time. default:30 seconds' )
    optparser.add_option( '-T' , '--timeout' , dest='time_out' , type='int', default='2', help='request time out. default:2 seconds' )
    (options , url) = optparser.parse_args()
   
    if len(url) == 0:
        print 'test url required'
        exit()
    
    signal.signal(signal.SIGINT, safe_exit)

    do_request(url[0], options.client_num, options.req_time, options.time_out)