#!/usr/bin/python
import requests, copy, json, datetime, os

REFLECTOR_ID_FILE="/root/reflector_id"
INITIATOR_ID_FILE="/root/initiator_id"

##############################################################################################################################################################3

start_reflector_template={u'agent': {u'primary': {u'config': u'primaryConfig',
                                                  u'id': u'vpma1',
                                                  u'type': u'inventory'}},
                          u'config': [{u'primaryConfig': {u'ietf-lmap-control:lmap': {u'schedules': {u'schedule': [{u'action': [{u'name': u'TWAMP-REFLECTOR-PM-PROBE0-action',
                                                                                                                                 u'option': [{u'id': u'source-ip',
                                                                                                                                              u'value': u'10.0.0.0'},
                                                                                                                                             {u'id': u'subnet',
                                                                                                                                              u'value': u'255.255.255.0'},
                                                                                                                                             {u'id': u'gateway',
                                                                                                                                              u'value': u'0.0.0.0'},
                                                                                                                                             {u'id': u'twamp',
                                                                                                                                              u'value': u'{"ietf-twamp:twamp":{"twamp-server":{"server-admin-state":true,"server-tcp-port":"862","dscp":0,"mode":"bits.unauthenticated"}}}'}],
                                                                                                                                 u'task': u'twamp-reflector'}],
                                                                                                                    u'name': u'TWAMP-REFLECTOR-PM-PROBE0',
                                                                                                                    u'start': u'immediate'}]}}}}],
                          u'test': {}}

##############################################################################################################################################################3

start_initiator_template=u'{"agent":{"primary":{"type":"inventory","id":"%(initiator_name)s","config":"primaryConfig"}},"test":{"extended-config":{"test":{"nameOnProbe":"twamp-initiator","selected":true,"testType":"twamp","testTypeLabel":"TWAMP Initiator"},"testPoints":[{"ip":"%(initiator_ip)s","model":"vPMA","name":"%(initiator_name)s","port":"80","selected":true,"tag":"Test Point A","type":"inventory","vendor":"viavi","version":"1.0"},{"model":"vPMA","name":"Test Point B","selected":false,"tag":"Test Point B","type":"inventory","vendor":"viavi","version":"1.0"}],"pluginConfig":{"setup":{"version":"2","schedMode":"immediate","periodicDuration":"00:05:00","periodicInterval":"00:15:00","pmInterval":"1","enableRandomStart":"false","randomOffset":"00:10:00","enableRt":"true","rtInterval":"%(interval)s","srcIPAddress":"%(initiator_src_ip)s","srcNetmask":"255.255.255.0","gatewayIPAddress":"192.168.2.1","srcStartingPort":"50000","configToServicesMap":{},"serviceParamToConfigMap":{}},"connections":{"version":"2","connections":[{"version":"2","name":"My Session","enabled":"true","source":null,"destination":{"version":"2","ipv4DestMgmtAddr":"%(reflector_src_ip)s","dstMgmtPort":"862"},"network":null,"numFlows":1,"flows":[{"version":"1","flowName":"Flow 1","flowInUse":"true","flowIpv4DiffServVal":"0","flowPad":"483","flowUdpPort":"50000","flowTxInterval":"10","flowReflWaitTimeout":"2"}]}]}},"testInfo":{"comments":"","name":"Lorna-%(test_name)s"}},"test-list-data":{"customer":{}}},"config":[{"primaryConfig":{"ietf-lmap-control:lmap":{"schedules":{"schedule":[{"name":"Lorna-%(test_name)s","start":"immediate","action":[{"name":"Lorna-%(test_name)s-action","task":"twamp-initiator","option":[{"id":"subnet","value":"255.255.255.0"},{"id":"gateway","value":"0.0.0.0"},{"id":"twamp","value":"{\\"ietf-twamp:twamp\\":{\\"twamp-client\\":{\\"client-admin-state\\":true,\\"mode-preference-chain\\":[{\\"priority\\":1,\\"mode\\":\\"bits.unauthenticated\\"}],\\"twamp-client-ctrl-connection\\":[{\\"ctrl-connection-name\\":\\"My Session\\",\\"client-ip\\":\\"%(initiator_src_ip)s\\",\\"server-ip\\":\\"%(reflector_src_ip)s\\",\\"server-tcp-port\\":862,\\"dscp\\":0,\\"twamp-session-request\\":[{\\"test-session-name\\":\\"Flow 1\\",\\"sender-ip\\":\\"%(initiator_src_ip)s\\",\\"sender-udp-port\\":\\"50000\\",\\"reflector-ip\\":\\"%(reflector_src_ip)s\\",\\"reflector-udp-port\\":\\"50000\\",\\"timeout\\":\\"2\\",\\"padding-length\\":\\"483\\",\\"dscp\\":\\"0\\"}]}],\\"twamp-session-sender\\":{\\"session-sender-admin-state\\":true,\\"twamp-sender-test-session\\":[{\\"test-session-name\\":\\"Flow 1\\",\\"fill-mode\\":\\"fill-mode.zero\\",\\"periodic-interval\\":\\"10\\",\\"periodic-interval-units\\":\\"milliseconds\\"}]}}}}"}],"destination":["Lorna-%(test_name)s-results-periodic-pm","Lorna-%(test_name)s-results-periodic-rt"]}]},{"name":"Lorna-%(test_name)s-results-periodic-pm","start":"Lorna-%(test_name)s-results-periodic-pm-event","action":[{"name":"Lorna-%(test_name)s-results-periodic-pm-action","task":"results","option":[{"id":"result-type","value":"pm"},{"id":"channel","value":"http://%(controller_ip)s:26088/restconf/operations/ietf-lmap-report:report"}]}]},{"name":"Lorna-%(test_name)s-results-periodic-rt","start":"Lorna-%(test_name)s-results-periodic-rt-event","action":[{"name":"Lorna-%(test_name)s-results-periodic-rt-action","task":"results","option":[{"id":"result-type","value":"rt"},{"id":"channel","value":"http://%(controller_ip)s:26088/restconf/operations/ietf-lmap-report:report"}]}]}]},"events":{"event":[{"name":"Lorna-%(test_name)s-results-periodic-pm-event","periodic":{"interval":60}},{"name":"Lorna-%(test_name)s-results-periodic-rt-event","periodic":{"interval":5}}]}}}}]}'

class FusionController:
    def __init__(self,inventory_ip):
        self.ip = inventory_ip
        self.i_url="http://%s:30100"%self.ip # inventory url
        self.c_url="http://%s:8080"%self.ip  # controller url
        self.h = {"Content-Type":"application/json"}
    def get_ma(self,ma_ip):
        inventory = requests.post("%s/vibe/inventory/v2/queryDevices"%self.i_url, headers=self.h).json()
        return [ma for ma in inventory if ma_ip in ma['connections'][0]['uri']][0]
        
    def start_reflector(self, reflector_ip, reflector_source_ip):
        ma = self.get_ma(reflector_ip)
        data = copy.deepcopy(start_reflector_template)
        data['agent']['primary']['id']=ma['name']
        config_options=data['config'][0]['primaryConfig']['ietf-lmap-control:lmap']['schedules']['schedule'][0]['action'][0]['option']
        config_options[0]['value']=reflector_source_ip
        r=requests.post("%s/sbi/restconf/data/ietf-lmap-control:lmap"%self.c_url,headers=self.h,data=json.dumps(data))
        open(REFLECTOR_ID_FILE,"w").write(r.json()["id"])

    def start_initiator(self, initiator_ip, initiator_source_ip, reflector_ip, reflector_source_ip, interval="1"):
        reflector = self.get_ma(reflector_ip)
        initiator = self.get_ma(initiator_ip)
        
	now = datetime.datetime.now()
	timestamp = '{}-{}-{}-{}:{}'.format(now.year, now.month, now.day, now.hour, now.minute)
        config_options={"initiator_name":initiator["name"],"initiator_ip":initiator_ip,"initiator_src_ip":initiator_source_ip,
                        "reflector_name":reflector["name"],"reflector_ip":reflector_ip,"reflector_src_ip":reflector_source_ip,
                        "interval":interval,"test_name":timestamp,"controller_ip":self.ip}
        r=requests.post("%s/sbi/restconf/data/ietf-lmap-control:lmap"%self.c_url,headers=self.h,data=start_initiator_template%config_options)
        open(INITIATOR_ID_FILE,"w").write(r.json()["id"])

    def stop_test(self, id):
       print("%s/sbi/restconf/operations/action/lmap?scheduleid=%s"%(self.c_url,id))
       requests.post("%s/sbi/restconf/operations/action/lmap?scheduleid=%s"%(self.c_url,id),headers=self.h,data=json.dumps({'action':'cancel'}))


def start_reflector(options):
    f = FusionController(options.controller_ip)
    f.start_reflector(options.reflector_ip,options.reflector_src_ip)

def start_initiator(options):
    f = FusionController(options.controller_ip)
    f.start_initiator(options.initiator_ip,options.initiator_src_ip,options.reflector_ip,options.reflector_src_ip,options.results_interval)

def stop_test(options,id_file):
    if os.path.exists(id_file):
        f = FusionController(options.controller_ip)
        f.stop_test(open(id_file).read())
        os.remove(id_file)

if __name__=="__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--controller-ip', help='IP of the Server', default='172.27.0.3')
    parser.add_argument('--results-interval', help='Results Interval', default='5')
    parser.add_argument('--reflector-ip', help='IP of the Agent to use as Reflector', default='10.0.6.4')
    parser.add_argument('--reflector-src-ip', help='Source IP of the reflector', default='10.0.7.3')
    parser.add_argument('--initiator-ip', help='IP of the Agent to use as initiator', default='10.0.6.5')
    parser.add_argument('--initiator-src-ip', help='Source IP of the reflector', default='10.0.7.4')
    parser.add_argument("--start-reflector", help="start reflector", action="store_true")
    parser.add_argument("--start-initiator", help="start initiator", action="store_true")
    parser.add_argument("--stop-reflector", help="stop reflector", action="store_true")
    parser.add_argument("--stop-initiator", help="stop initiator", action="store_true")

    options = parser.parse_args()
    if options.start_reflector:
        start_reflector(options)
    if options.start_initiator:
        start_initiator(options)
    if options.stop_reflector:
        stop_test(options,REFLECTOR_ID_FILE)
    if options.stop_initiator:
        stop_test(options,INITIATOR_ID_FILE)