# coding: utf-8

# Adapted from code in https://github.com/konomae/stunpy

import binascii
import random
import socket
import struct

STUN_PORT = 3478

FAMILY_IPv4 = b"\x01"
FAMILY_IPv6 = b"\x02"

# 0b00: Request
# 0b01: Binding
BINDING_REQUEST_SIGN = b"\x00\x01"  # 16bit (2bytes)

BINDING_RESPONSE_ERROR = b"\x01\x11"
BINDING_RESPONSE_SUCCESS = b"\x01\x01"
MAGIC_COOKIE = b"\x21\x12\xA4\x42"  # 32bit (4bytes)

CHANGE_RESPONSE_SUCCESS = b"\x01\x03"

CONNECT_REQUEST = b"\x00\x0a"
CONNECT_RESPONSE_SUCCESS = b"\x01\x0a"

CONNECTION_BIND_REQUEST = b"\x00\x0b"
CONNECTION_BIND_SUCCESS = b"\x01\x0b"

CONNECTION_REFRESH_REQUEST = b"\x00\x04"

# STUN Attribute Registry
MAPPED_ADDRESS = b"\x00\x01"
RESPONSE_ADDRESS = b"\x00\x02"
CHANGE_REQUEST = b"\x00\x03"
SOURCE_ADDRESS = b"\x00\x04"
CHANGED_ADDRESS = b"\x00\x05"
USERNAME = b"\x00\x06"
PASSWORD = b"\x00\x07"
MESSAGE_INTEGRITY = b"\x00\x08"
ERROR_CODE = b"\x00\x09"
UNKNOWN_ATTRIBUTES = b"\x00\x0A"
REFLECTED_FROM = b"\x00\x0B"
REALM = b"\x00\x14"
NONCE = b"\x00\x15"
XOR_MAPPED_ADDRESS = b"\x80\x82"
XOR_PEER_ADDRESS = b"\x00\x12"
REQUESTED_TRANSPORT = b"\x00\x19"
TRANSPORT_TCP = b"\x06"
LIFETIME = b"\x00\x0d"
REQUESTED_ADDRESS_TYPE = b"\x00\x17"
RESERVED3 = b"\x00\x00\x00"
RESERVED1 = b"\x00"
CONNECTION_ID = b"\x00\x2a"

STUN_ATTRIBUTE_NAMES = {
    MAPPED_ADDRESS: "MAPPED-ADDRESS",
    RESPONSE_ADDRESS: "RESPONSE-ADDRESS",
    CHANGE_REQUEST: "CHANGE-REQUEST",
    SOURCE_ADDRESS: "SOURCE-ADDRESS",
    CHANGED_ADDRESS: "CHANGED-ADDRESS",
    USERNAME: "USERNAME",
    PASSWORD: "PASSWORD",
    MESSAGE_INTEGRITY: "MESSAGE-INTEGRITY",
    ERROR_CODE: "ERROR-CODE",
    UNKNOWN_ATTRIBUTES: "UNKNOWN-ATTRIBUTES",
    REFLECTED_FROM: "REFLECTED-FROM",
    REALM: "REALM",
    NONCE: "NONCE",
    XOR_MAPPED_ADDRESS: "XOR-MAPPED-ADDRESS",
    b"\x00\x20": "XOR-MAPPED-ADDRESS",
}


def generate_transaction_id():
    tid = []
    for i in range(24):  # 96bits (12bytes)
        tid.append(random.choice("0123456789ABCDEF"))
    return binascii.a2b_hex("".join(tid))


def build_binding_request(transaction_id):
    if len(transaction_id) != 12:
        raise Exception("Invalid transaction id")

    body_length = b"\x00\x00"
    return b"".join([BINDING_REQUEST_SIGN, body_length, MAGIC_COOKIE, transaction_id])


def build_change_request(transaction_id):
    if len(transaction_id) != 12:
        raise Exception("Invalid transaction id")

    body_length = b"\x00\x18"
    return b"".join(
        [
            CHANGE_REQUEST,
            body_length,
            MAGIC_COOKIE,
            transaction_id,
            REQUESTED_TRANSPORT,
            b"\x00\x04",
            TRANSPORT_TCP,
            RESERVED3,
            LIFETIME,
            b"\x00\x04",
            b"\x00\x00\x02\x58",
            REQUESTED_ADDRESS_TYPE,
            b"\x00\x04",
            FAMILY_IPv4,
            RESERVED3,
        ]
    )


def build_connection_bind_request(transaction_id, connection_id):
    if len(transaction_id) != 12:
        raise Exception("Invalid transaction id")

    if len(connection_id) != 4:
        raise Exception("Invalid connection id {}".format(len(connection_id)))

    body_length = b"\x00\x08"
    return b"".join(
        [
            CONNECTION_BIND_REQUEST,
            body_length,
            MAGIC_COOKIE,
            transaction_id,
            CONNECTION_ID,
            b"\x00\x04",
            connection_id,
        ]
    )


def build_connection_refresh_request(transaction_id):
    if len(transaction_id) != 12:
        raise Exception("Invalid transaction id")

    body_length = b"\x00\x08"
    return b"".join(
        [
            CONNECTION_REFRESH_REQUEST,
            body_length,
            MAGIC_COOKIE,
            transaction_id,
            LIFETIME,
            b"\x00\x04",
            b"\x00\x00\x02\x58",
        ]
    )


def build_connect_request(transaction_id, xoraddr=None):
    if len(transaction_id) != 12:
        raise Exception("Invalid transaction id")

    if xoraddr is None:
        raise Exception("Invalid connect address")

    body_length = b"\x00\x0c"
    return b"".join(
        [
            CONNECT_REQUEST,
            body_length,
            MAGIC_COOKIE,
            transaction_id,
            XOR_PEER_ADDRESS,
            b"\x00\x08",
            xoraddr,
        ]
    )


def validate_response(buf, transaction_id):
    if not buf or len(buf) < 20:
        raise Exception("Response too short")

    response_magic_cookie = buf[4:8]
    if MAGIC_COOKIE != response_magic_cookie:
        raise Exception("Invalid magic cookie")

    response_transaction_id = buf[8:20]
    if transaction_id != response_transaction_id:
        raise Exception("invalid transaction id")


def ip_to_bytes(ip, xor):
    octets = [binascii.a2b_hex("%02x" % int(o)) for o in ip.split(".")]
    addr_int = struct.unpack("!I", b"".join(octets))[0]

    if xor:
        magicCookieBytesInt = int(binascii.b2a_hex(MAGIC_COOKIE), 16)
        addr_int = magicCookieBytesInt ^ addr_int

    addr_bytes = binascii.a2b_hex("%08x" % addr_int)
    return addr_bytes


def port_to_bytes(port, xor):
    if xor:
        magicCookieHighBytesInt = int(binascii.b2a_hex(MAGIC_COOKIE[:2]), 16)
        port = magicCookieHighBytesInt ^ port

    port_bytes = binascii.a2b_hex("%04x" % port)
    return port_bytes


def read_mapped_address(attr_type, attr_body, attr_len):
    assert attr_type in (MAPPED_ADDRESS, XOR_MAPPED_ADDRESS, b"\x00\x20")
    assert attr_body[:1] == b"\x00"

    family_bytes = attr_body[1:2]
    port_bytes = attr_body[2:4]
    addr_bytes = attr_body[4:attr_len]
    xor = attr_type in (XOR_MAPPED_ADDRESS, b"\x00\x20")

    family_text = ""
    assert family_bytes in (FAMILY_IPv4, FAMILY_IPv6)
    if family_bytes == FAMILY_IPv4:
        family_text = "IPv4"
    elif family_bytes == FAMILY_IPv6:
        family_text = "IPv6"

    # TODO: IPv6s
    port = int(binascii.b2a_hex(port_bytes), 16)
    addr_int = int(binascii.b2a_hex(addr_bytes), 16)

    if xor:
        # port
        magicCookieHighBytesInt = int(binascii.b2a_hex(MAGIC_COOKIE[:2]), 16)
        port = magicCookieHighBytesInt ^ port

        # addr
        magicCookieBytesInt = int(binascii.b2a_hex(MAGIC_COOKIE), 16)
        addr_int = magicCookieBytesInt ^ addr_int

    octets = struct.pack("!I", addr_int)
    ip = ".".join([str(c) for c in octets])
    return dict(
        name=STUN_ATTRIBUTE_NAMES[attr_type], ip=ip, port=port, family=family_text
    )


def read_attributes(attributes, body_length):
    pos = 0
    parsed_attributes = []
    while pos < body_length:
        attr_type = attributes[pos : pos + 2]  # 16bit (2bytes)
        attr_len = int(
            binascii.b2a_hex(attributes[pos + 2 : pos + 4]), 16
        )  # 16bit (2bytes)
        attr_body = attributes[pos + 4 : pos + 4 + attr_len]

        if attr_type in (MAPPED_ADDRESS, XOR_MAPPED_ADDRESS, b"\x00\x20"):
            parsed_attributes.append(
                read_mapped_address(attr_type, attr_body, attr_len)
            )
        else:
            if attr_type == b"\x80\x22":
                attr_body = attr_body.decode("utf-8")
            else:
                attr_body = attr_body.hex()
            parsed_attributes.append(
                dict(
                    name=STUN_ATTRIBUTE_NAMES.get(attr_type),
                    attr_type=attr_type.hex(),
                    attr_body=attr_body,
                    attr_len=attr_len,
                )
            )

        attr_head = 2 + 2  # attr_type + attr_len
        remain = attr_len % 4
        padding = 4 - remain if remain else 0
        pos += attr_head + attr_len + padding

    return parsed_attributes


def is_error(body):
    for element in body:
        if element["name"] == "ERROR-CODE":
            return True

    return False


def get_error(body):
    for element in body:
        if element["name"] == "ERROR-CODE":
            body = binascii.unhexlify(element["attr_body"])
            return {"class": body[:3], "code": body[3], "reason": body[4:]}

    return False


class StunClient(object):
    def __init__(self, host, port=STUN_PORT):
        self.sock = None
        self.transaction_id = None
        self.req = None
        self.transaction_id = generate_transaction_id()

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("0.0.0.0", 0))
        self.host = host
        self.port = port
        self.sock.connect((host, port))

    def send_tcp_change_request(self):
        self.req = build_change_request(self.transaction_id)
        self.sock.send(self.req)

    def send_binding_request(self):
        self.req = build_binding_request(self.transaction_id)
        self.sock.send(self.req)

    def send_connect_request(self, host=None, port=None, xoraddr=None):
        if host is not None and port is not None:
            xorhost = ip_to_bytes(host, True)
            xorport = port_to_bytes(port, True)
            addr = b"".join([xorport, xorhost])
        elif xoraddr is not None:
            addr = xoraddr
        else:
            raise Exception("Invalid arguments")

        self.req = build_connect_request(self.transaction_id, addr)
        self.sock.send(self.req)

    def send_connection_bind_request(self, connection_id):
        self.req = build_connection_bind_request(self.transaction_id, connection_id)
        self.sock.send(self.req)

    def receive_response(self):
        buf = self.sock.recv(2048)
        validate_response(buf, self.transaction_id)

        body_length = int(binascii.b2a_hex(buf[2:4]), 16)
        attributes = buf[20:]
        assert len(attributes) == body_length

        return read_attributes(attributes, body_length)

    def send_refresh_request(self):
        self.req = build_connection_refresh_request(self.transaction_id)
        self.sock.send(self.req)

    def get_socket(self):
        return self.sock

    def close(self):
        if self.sock:
            self.sock.close()
