#!/usr/bin/env python

'''Generate multiple RTP streams to test how mediaproxy copes with load

Copyright 2003-2005 Dan Pascu
'''

import asyncore
import socket
import sys, os, errno, traceback
import struct
import re
from time import time

## Script settings

controlSock      = '/var/run/mediaproxy.sock'

generatorIP = '127.0.0.1'
minPort = 10000
maxPort = 30000
TOS = 0xb8
dataSize = 172
packetDelay = 0.02 ## 20 ms
sendAtOnce = 5     ## How many packets to send at once in burst mode
maxDataSize = 16384

dummyData = dataSize*"0"

## Helper functions

def warning(msg):
    print >>sys.stderr, "warning: " + str(msg)

def error(msg):
    print >>sys.stderr, "error: " + str(msg)

## Classes

class GeneratorError(Exception):
    pass

class RTPGenerator(asyncore.dispatcher):
    def __init__(self, mode, session=None, map=None):
        sock = self.__findFreeSocket()
        asyncore.dispatcher.__init__(self, sock, map)
        self.addr = sock.getsockname()
        self.ip = self.addr[0]
        self.port = self.addr[1]
        self.domain = "dummydomain"
        self.local = "remote"
        self.session = session or ("%s:%s" % self.addr)
        self.mode = mode
        command = "%(mode)s %(session)s %(ip)s:%(port)d:audio %(ip)s dummy remote dummy remote dummy-agent info=from:test@test.com,to:test@test.com\n" % self.__dict__
        proxy = self.sendCommand(command).split()
        if proxy:
            self.proxy = (proxy[0], int(proxy[1]))
        else:
            raise GeneratorError, "No answer from mediaproxy"

    def sendCommand(self, command):
        if type(controlSock) == type(()):
            socktype = socket.AF_INET
        else:
            socktype = socket.AF_UNIX
        s = socket.socket(socktype, socket.SOCK_STREAM)
        try:
            s.connect(controlSock)
            s.send(command)
            return s.recv(maxDataSize)
        finally:
            s.close()
    
    def end(self):
        self.sendCommand("delete %s info=\n" % self.session)

    def writable(self):
        return 0

    def handle_read(self):
        data = self.recv(65536)

    def handle_write(self):
        try:
            for x in range(0, sendAtOnce):
                self.socket.sendto(dummyData, self.proxy)
        except KeyboardInterrupt:
            raise asyncore.ExitNow
        except socket.error, why:
            if why[0] != errno.EWOULDBLOCK:
                raise socket.error, why

    def handle_close(self):
        self.close()
        self.end()

    def __findFreeSocket(self):
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            sock.setsockopt(socket.IPPROTO_IP, socket.IP_TOS, TOS)
        except socket.error:
            pass ## Not critical
        for port in range(minPort, maxPort, 2):
            try:
                try:
                    sock.bind((generatorIP, port))
                except socket.error, why:
                    if why[0] in (errno.EADDRINUSE, errno.EACCES):
                        continue
                    else:
                        raise GeneratorError, why[1]
                else:
                    return sock
            # maybe a separate check for 'except GeneratorError' too
            except:
                ## Print traceback
                info = sys.exc_info()
                try:
                    traceback.print_exception(*info)
                finally:
                    del info
                break
        sock.close()
        return None


def loop(map=None):
    if map is None:
        map = asyncore.socket_map
    delay = packetDelay * sendAtOnce
    next = time() + delay
    while map:
        now = time()
        if now >= next:
            for dev in map.values():
                dev.handle_write()
            next += delay
            timeout = next - time()
            if timeout < 0.0:
                timeout = 0.0
        else:
            timeout = next - now
        asyncore.poll(timeout, map)


if __name__ == '__main__':
    sys.path[1:1] = [os.path.realpath((os.path.dirname(sys.argv[0]) or '.') + '/../')]
    sys.path[1:1] = [os.path.realpath((os.path.dirname(sys.argv[0]) or '.') + '/../modules')]

    from optparse import OptionParser
    from options import ExtendedOptions

    version = '%d.%d' % sys.version_info[0:2]
    
    parser = OptionParser(option_class=ExtendedOptions, version="%prog 1.0")
    parser.add_option("--ip", dest="ip", type="ip", default="127.0.0.1",
                      help="IP used to send/receive the RTP streams (127.0.0.1)",
                      metavar="IP")
    parser.add_option("--proxy", dest="controlSock", default='/var/run/mediaproxy.sock',
                      help="mediaproxy control socket (/var/run/mediaproxy.sock)",
                      type="controlport", metavar="Host|File")
    parser.add_option("--g711", action="store_true", dest="g711", default=0,
                      help="create a load equivalent of a G711 codec")
    parser.add_option("--g729", action="store_true", dest="g729", default=0,
                      help="create a load equivalent of a G729 codec")
    parser.add_option("--g723", action="store_true", dest="g723", default=0,
                      help="create a load equivalent of a G723 codec")
    parser.add_option("--gsm", action="store_true", dest="gsm", default=0,
                      help="create a load equivalent of a GSM codec")
    parser.add_option("--burst", action="store_true", dest="burst", default=0,
                      help="send packets in bursts instead of evenly spaced")
    parser.add_option("--rate", dest="rate", default=64, type="int",
                      help="data rate in kpbs 0-512 (64)",
                      metavar="Rate")
    parser.add_option("--count", dest="count", default=1, type="int",
                      help="simulate this many streams of the given rate 1-50 (1)",
                      metavar="Count")
    (options, args) = parser.parse_args()
    controlSock = options.controlSock
    generatorIP = options.ip
    if options.rate<0 or options.rate>512:
        print "Rate must be between 0 and 512 kbps"
        sys.exit(1)
    if options.count<1 or options.count>50:
        print "Count must be between 1 and 50"
        sys.exit(1)
    if not options.burst:
        sendAtOnce = 1
    if options.g711:
        dataSize = 172
        payload = '\x00'
    elif options.g729:
        dataSize = 32
        payload = '\x12'
    elif options.g723:
        dataSize = 2
        payload = '\x04'
    elif options.gsm:
        dataSize = 45
        payload = '\x03'
    else:
        dataSize = int((options.rate * 1024.0 / 8) * packetDelay) # -28 for ip/udp, -42 for all headers
        payload = '\x7f'
    dummyData = '\x00' + payload + (dataSize-2)*"\x00"
    senders = []
    try:
        for i in range(options.count):
            sender = RTPGenerator('request')
            receiver = RTPGenerator('lookup', sender.session)
            senders.append(sender)
        loop()
    except GeneratorError, why:
        print >>sys.stderr, "Generator error: %s." % why
    except (KeyboardInterrupt, asyncore.ExitNow):
        print >>sys.stderr, "Got keyboard interrupt. Exiting."
    except socket.error, why:
        print >>sys.stderr, "Cannot connect to proxy: %s" % why[1]
    [sender.end() for sender in senders]

