# retoor <retoor@molodetz.nl>

import asyncio
import base64

from devplacepy.services.nodes import tunnel as tunnel_module
from devplacepy.services.nodes.policy import DestinationPolicy
from devplacepy.services.nodes.tunnel import Relay, TunnelBridge
from tests.conftest import run_async

TARGET = DestinationPolicy().check("https://www.reddit.com/r/python.json")


class Donor:
    def __init__(self, bridge, accept=True, garbage=False):
        self.bridge = bridge
        self.accept = accept
        self.garbage = garbage
        self.tunnels = {}
        self.sent = []

    async def send(self, frame_type, **fields):
        self.sent.append({"type": frame_type, **fields})
        if frame_type == "tunnel_open":
            self.bridge.on_frame(self, {"type": "tunnel_opened", "conn_id": fields["conn_id"], "ok": self.accept})
        elif frame_type == "tunnel_data":
            raw = base64.b64decode(fields["b64"])
            reply = "***not base64***" if self.garbage else base64.b64encode(raw.upper()).decode()
            self.bridge.on_frame(self, {"type": "tunnel_data", "conn_id": fields["conn_id"], "b64": reply})


async def session_through(donor, connect_line, payload=b"hello"):
    relay = Relay(donor, TARGET)
    server = await asyncio.start_server(relay.handle, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    reader, writer = await asyncio.open_connection("127.0.0.1", port)
    writer.write(connect_line.encode() + b"\r\n\r\n")
    await writer.drain()
    status = await asyncio.wait_for(reader.readline(), 5)
    echoed = b""
    if b" 200 " in status:
        await reader.readuntil(b"\r\n")
        writer.write(payload)
        await writer.drain()
        try:
            echoed = await asyncio.wait_for(reader.read(len(payload)), 5)
        except asyncio.TimeoutError:
            echoed = b""
    writer.close()
    server.close()
    await asyncio.sleep(0.05)
    return status, echoed, relay


def go(donor, line, payload=b"hello"):
    return run_async(session_through(donor, line, payload))


def test_connect_to_the_target_relays_both_ways():
    donor = Donor(TunnelBridge())
    status, echoed, relay = go(donor, "CONNECT www.reddit.com:443 HTTP/1.1")
    assert b" 200 " in status and echoed == b"HELLO"
    assert [item["type"] for item in donor.sent][:2] == ["tunnel_open", "tunnel_data"]
    assert donor.sent[0]["host"] == "www.reddit.com" and donor.sent[0]["port"] == 443
    assert relay.up == 5 and relay.down == 5


def test_any_other_host_or_port_is_refused_without_bothering_the_donor():
    for line in ("CONNECT evil.example:443 HTTP/1.1", "CONNECT www.reddit.com:22 HTTP/1.1", "CONNECT 10.0.0.1:443 HTTP/1.1"):
        donor = Donor(TunnelBridge())
        status, _, relay = go(donor, line)
        assert b" 403 " in status and donor.sent == []
        assert relay.error == "host_not_allowed"


def test_only_connect_is_accepted():
    donor = Donor(TunnelBridge())
    status, _, _ = go(donor, "GET http://www.reddit.com/ HTTP/1.1")
    assert b" 405 " in status and donor.sent == []


def test_a_donor_that_refuses_is_a_bad_gateway():
    donor = Donor(TunnelBridge(), accept=False)
    status, _, relay = go(donor, "CONNECT www.reddit.com:443 HTTP/1.1")
    assert b" 502 " in status and relay.error == "tunnel_refused"


def test_the_relay_has_a_byte_cap(monkeypatch):
    monkeypatch.setattr(tunnel_module, "RELAY_CAP_BYTES", 8)
    donor = Donor(TunnelBridge())
    _, echoed, relay = go(donor, "CONNECT www.reddit.com:443 HTTP/1.1", b"0123456789abcdef")
    assert relay.error == "body_too_large" and len(echoed) < 16


def test_invalid_base64_from_the_donor_ends_the_tunnel():
    donor = Donor(TunnelBridge(), garbage=True)
    _, echoed, relay = go(donor, "CONNECT www.reddit.com:443 HTTP/1.1")
    assert relay.error == "bad_base64" and echoed == b""


def test_releasing_a_session_closes_its_relays():
    bridge = TunnelBridge()
    donor = Donor(bridge)
    relay = Relay(donor, TARGET)
    donor.tunnels[relay.conn_id] = relay
    bridge.release_session(donor)
    assert donor.tunnels == {} and relay.frames.get_nowait() is tunnel_module.CLOSED


def test_frames_for_unknown_connections_are_not_consumed():
    bridge = TunnelBridge()
    assert bridge.on_frame(Donor(bridge), {"type": "tunnel_data", "conn_id": "nope", "b64": ""}) is False