#!/usr/bin/env python3

import aiohttp
import asyncio
import json
import requests
import random
import argparse
import string
import time

from pwn import log


def random_string(string_length=10):
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for i in range(string_length))


def transaction_id(x=10):
    return random_string(x)


class Exploit:
    def __init__(self, *args, **kwargs):
        pass

    def exploit(self, *args, tasks=None, **kwargs):
        if tasks is None:
            return
        loop = asyncio.get_event_loop()
        tasks_ = []
        for task, values in tasks:
            if len(values) > 0:
                tasks_.append(task(*values))
            else:
                tasks_.append(task())
        wait_tasks = asyncio.wait(tasks_)

        loop.run_until_complete(
            wait_tasks
        )

        loop.close()

    def _class_name(self):
        return self.__class__.__name__

    def log_success(self, msg):
        log.success('{}: {}'.format(self._class_name(), msg))

    def log_info(self, msg):
        log.info('{}: {}'.format(self._class_name(), msg))

    def log_error(self, msg):
        log.error('{}: {}'.format(self._class_name(), msg))

    @staticmethod
    def parse_args():
        parser = argparse.ArgumentParser(description='Jira tests')
        parser.add_argument('-s', '--server', type=str,
            default='http://localhost', help='Janus server address')
        parser.add_argument('-p', '--port', type=int,
            default=8088, help='Janus server port')

        args = parser.parse_args()
        server, port = args.server, args.port
        return server, port


class JanusPlugin:
    def __init__(self, session, url):
        self._queue = asyncio.Queue()
        self._session = session
        self._url = url

    async def send(self, payload, bypass=False):
        message = {"janus": "message", "transaction": transaction_id()}
        message.update(payload)
        async with self._session._http.post(self._url, json=message) as response:
            data = await response.json()
            log.info('post[response]: {}'.format(data))
            
            if bypass:
                return None
                
            if data["janus"] != "ack":
                log.warn('data["janus"] != "ack --> something is wrong"')
                return data

        response = await self._queue.get()
        assert response["transaction"] == message["transaction"]
        return response


class JanusSession:
    def __init__(self, url, name):
        self._http = None
        self._poll_task = None
        self._plugins = {}
        self._root_url = url
        self._session_url = None
        self._jsep_answer = None
        self._jsep_offer  = None
        self.name = name

    async def attach(self, plugin_name: str) -> JanusPlugin:
        message = {
            "janus": "attach",
            "plugin": plugin_name,
            "transaction": transaction_id(),
        }
        async with self._http.post(self._session_url, json=message) as response:
            data = await response.json()
            assert data["janus"] == "success"
            plugin_id = data["data"]["id"]
            plugin = JanusPlugin(self, self._session_url + "/" + str(plugin_id))
            self._plugins[plugin_id] = plugin
            return plugin

    async def create(self):
        self._http = aiohttp.ClientSession()
        message = {"janus": "create", "transaction": transaction_id()}
        async with self._http.post(self._root_url, json=message) as response:
            data = await response.json()
            assert data["janus"] == "success"
            session_id = data["data"]["id"]
            self._session_url = self._root_url + "/" + str(session_id)

        self._poll_task = asyncio.ensure_future(self._poll())

    async def destroy(self):
        if self._poll_task:
            self._poll_task.cancel()
            self._poll_task = None

        if self._session_url:
            message = {"janus": "destroy", "transaction": transaction_id()}
            async with self._http.post(self._session_url, json=message) as response:
                data = await response.json()
                assert data["janus"] == "success"
            self._session_url = None

        if self._http:
            await self._http.close()
            self._http = None

    async def _poll(self):
        while True:
            params = {
                "maxev": 1,
                "rid": int(time.time() * 1000)
            }
            
            async with self._http.get(self._session_url, params=params) as response:
                data = await response.json()

                log.info("poller: name: {0}, data: {1}\n".format(self.name, data))

                if data["janus"] == "event":
                    plugin = self._plugins.get(data["sender"], None)
                    if plugin:
                        await plugin._queue.put(data)

                    if "jsep" in data:
                        if data["jsep"]["type"] == "answer":
                            self._jsep_answer = data["jsep"]
                        elif data["jsep"]["type"] == "offer":
                            self._jsep_offer = data["jsep"]

    async def get_answer(self):
        answer = None
        while not answer:
            await asyncio.sleep(1)
            answer = self._jsep_answer
        return answer

    async def get_offer(self):
        offer = None
        while not offer:
            await asyncio.sleep(1)
            offer = self._jsep_offer
        return offer
