#!/usr/bin/env python
# License: GPLv3 Copyright: 2015, Kovid Goyal <kovid at kovidgoyal.net>


import os
import socket
import weakref
from collections import deque
from hashlib import sha1
from struct import error as struct_error
from struct import pack, unpack_from
from threading import Lock

from calibre import as_unicode
from calibre.srv.http_response import HTTPConnection, create_http_handler
from calibre.srv.loop import RDWR, READ, WRITE, Connection, HandleInterrupt, ServerLoop
from calibre.srv.utils import DESIRED_SEND_BUFFER_SIZE
from calibre.utils.speedups import ReadOnlyFileBuffer
from calibre_extensions.speedup import utf8_decode
from calibre_extensions.speedup import websocket_mask as fast_mask
from polyglot import http_client
from polyglot.binary import as_base64_unicode
from polyglot.queue import Empty, Queue

HANDSHAKE_STR = (
    "HTTP/1.1 101 Switching Protocols\r\n"
    "Upgrade: WebSocket\r\n"
    "Connection: Upgrade\r\n"
    "Sec-WebSocket-Accept: %s\r\n\r\n"
)
GUID_STR = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

CONTINUATION = 0x0
TEXT = 0x1
BINARY = 0x2
CLOSE = 0x8
PING = 0x9
PONG = 0xA
CONTROL_CODES = (CLOSE, PING, PONG)
ALL_CODES = CONTROL_CODES + (CONTINUATION, TEXT, BINARY)

CHUNK_SIZE = 16 * 1024
SEND_CHUNK_SIZE = DESIRED_SEND_BUFFER_SIZE - 16

NORMAL_CLOSE = 1000
SHUTTING_DOWN = 1001
PROTOCOL_ERROR = 1002
UNSUPPORTED_DATA = 1003
INCONSISTENT_DATA = 1007
POLICY_VIOLATION = 1008
MESSAGE_TOO_BIG = 1009
UNEXPECTED_ERROR = 1011

RESERVED_CLOSE_CODES = (1004,1005,1006,)


class ReadFrame:  # {{{

    def __init__(self):
        self.header_buf = bytearray(14)
        self.rbuf = bytearray(CHUNK_SIZE)
        self.empty = memoryview(b'')
        self.reset()

    def reset(self):
        self.header_view = memoryview(self.header_buf)[:6]
        self.state = self.read_header

    def __call__(self, conn):
        return self.state(conn)

    def read_header(self, conn):
        num_bytes = conn.recv_into(self.header_view)
        if num_bytes == 0:
            return
        read_bytes = 6 - len(self.header_view) + num_bytes
        if read_bytes > 2:
            b1, b2 = self.header_buf[0], self.header_buf[1]
            self.fin = bool(b1 & 0b10000000)
            if b1 & 0b01110000:
                conn.log.error('RSV bits set in frame from client')
                conn.websocket_close(PROTOCOL_ERROR, 'RSV bits set')
                return

            self.opcode = b1 & 0b1111
            self.is_control = self.opcode in CONTROL_CODES
            if self.opcode not in ALL_CODES:
                conn.log.error('Unknown OPCODE from client: %r' % self.opcode)
                conn.websocket_close(PROTOCOL_ERROR, 'Unknown OPCODE: %r' % self.opcode)
                return
            if not self.fin and self.is_control:
                conn.log.error('Fragmented control frame from client')
                conn.websocket_close(PROTOCOL_ERROR, 'Fragmented control frame')
                return

            mask = b2 & 0b10000000
            if not mask:
                conn.log.error('Unmasked packet from client')
                conn.websocket_close(PROTOCOL_ERROR, 'Unmasked packet not allowed')
                self.reset()
                return
            self.payload_length = l = b2 & 0b01111111
            if self.is_control and l > 125:
                conn.log.error('Too large control frame from client')
                conn.websocket_close(PROTOCOL_ERROR, 'Control frame too large')
                self.reset()
                return
            header_len = 6 + (0 if l < 126 else 2 if l == 126 else 8)
            if header_len <= read_bytes:
                self.process_header(conn)
            else:
                self.header_view = memoryview(self.header_buf)[read_bytes:header_len]
                self.state = self.finish_reading_header
        else:
            self.header_view = self.header_view[num_bytes:]

    def finish_reading_header(self, conn):
        num_bytes = conn.recv_into(self.header_view)
        if num_bytes == 0:
            return
        if num_bytes >= len(self.header_view):
            self.process_header(conn)
        else:
            self.header_view = self.header_view[num_bytes:]

    def process_header(self, conn):
        if self.payload_length < 126:
            self.mask = memoryview(self.header_buf)[2:6]
        elif self.payload_length == 126:
            self.payload_length, = unpack_from(b'!H', self.header_buf, 2)
            self.mask = memoryview(self.header_buf)[4:8]
        else:
            self.payload_length, = unpack_from(b'!Q', self.header_buf, 2)
            self.mask = memoryview(self.header_buf)[10:14]
        self.frame_starting = True
        self.bytes_received = 0
        if self.payload_length <= CHUNK_SIZE:
            if self.payload_length == 0:
                conn.ws_data_received(self.empty, self.opcode, True, True, self.fin)
                self.reset()
            else:
                self.rview = memoryview(self.rbuf)[:self.payload_length]
                self.state = self.read_packet
        else:
            self.rview = memoryview(self.rbuf)
            self.state = self.read_payload

    def read_packet(self, conn):
        num_bytes = conn.recv_into(self.rview)
        if num_bytes == 0:
            return
        if num_bytes >= len(self.rview):
            data = memoryview(self.rbuf)[:self.payload_length]
            fast_mask(data, self.mask)
            conn.ws_data_received(data, self.opcode, True, True, self.fin)
            self.reset()
        else:
            self.rview = self.rview[num_bytes:]

    def read_payload(self, conn):
        num_bytes = conn.recv_into(self.rview, min(len(self.rview), self.payload_length - self.bytes_received))
        if num_bytes == 0:
            return
        data = memoryview(self.rbuf)[:num_bytes]
        fast_mask(data, self.mask, self.bytes_received)
        self.bytes_received += num_bytes
        frame_finished = self.bytes_received >= self.payload_length
        conn.ws_data_received(data, self.opcode, self.frame_starting, frame_finished, self.fin)
        self.frame_starting = False
        if frame_finished:
            self.reset()

# }}}

# Sending frames {{{


def create_frame(fin, opcode, payload, mask=None, rsv=0):
    if isinstance(payload, str):
        payload = payload.encode('utf-8')
    l = len(payload)
    header_len = 2 + (0 if l < 126 else 2 if 126 <= l <= 65535 else 8) + (0 if mask is None else 4)
    frame = bytearray(header_len + l)
    if l > 0:
        frame[-l:] = payload
    frame[0] = (opcode & 0b1111) | (0b10000000 if fin else 0) | (rsv & 0b01110000)
    if l < 126:
        frame[1] = l
    elif 126 <= l <= 65535:
        frame[2:4] = pack(b'!H', l)
        frame[1] = 126
    else:
        frame[2:10] = pack(b'!Q', l)
        frame[1] = 127
    if mask is not None:
        frame[1] |= 0b10000000
        frame[header_len-4:header_len] = mask
        if l > 0:
            fast_mask(memoryview(frame)[-l:], mask)

    return memoryview(frame)


class MessageWriter:

    def __init__(self, buf, mask=None, chunk_size=None):
        self.buf, self.data_type, self.mask = buf, BINARY, mask
        if isinstance(buf, str):
            self.buf, self.data_type = ReadOnlyFileBuffer(buf.encode('utf-8')), TEXT
        elif isinstance(buf, bytes):
            self.buf = ReadOnlyFileBuffer(buf)
        buf = self.buf
        self.chunk_size = chunk_size or SEND_CHUNK_SIZE
        try:
            pos = buf.tell()
            buf.seek(0, os.SEEK_END)
            self.size = buf.tell() - pos
            buf.seek(pos)
        except Exception:
            self.size = None
        self.first_frame_created = self.exhausted = False

    def create_frame(self):
        if self.exhausted:
            return None
        buf = self.buf
        raw = buf.read(self.chunk_size)
        has_more = True if self.size is None else self.size > buf.tell()
        fin = 0 if has_more and raw else 1
        opcode = 0 if self.first_frame_created else self.data_type
        self.first_frame_created, self.exhausted = True, bool(fin)
        return ReadOnlyFileBuffer(create_frame(fin, opcode, raw, self.mask))
# }}}


conn_id = 0


class UTF8Decoder:  # {{{

    def __init__(self):
        self.reset()

    def __call__(self, data):
        ans, self.state, self.codep = utf8_decode(data, self.state, self.codep)
        return ans

    def reset(self):
        self.state = 0
        self.codep = 0
# }}}


class WebSocketConnection(HTTPConnection):

    # Internal API {{{
    in_websocket_mode = False
    websocket_handler = None

    def __init__(self, *args, **kwargs):
        global conn_id
        HTTPConnection.__init__(self, *args, **kwargs)
        self.sendq = Queue()
        self.control_frames = deque()
        self.cf_lock = Lock()
        self.sending = None
        self.send_buf = None
        self.frag_decoder = UTF8Decoder()
        self.ws_close_received = self.ws_close_sent = False
        conn_id += 1
        self.websocket_connection_id = conn_id
        self.stop_reading = False

    def finalize_headers(self, inheaders):
        upgrade = inheaders.get('Upgrade', '')
        key = inheaders.get('Sec-WebSocket-Key', None)
        conn = {x.strip().lower() for x in inheaders.get('Connection', '').split(',')}
        if key is None or upgrade.lower() != 'websocket' or 'upgrade' not in conn:
            return HTTPConnection.finalize_headers(self, inheaders)
        ver = inheaders.get('Sec-WebSocket-Version', 'Unknown')
        try:
            ver_ok = int(ver) >= 13
        except Exception:
            ver_ok = False
        if not ver_ok:
            return self.simple_response(http_client.BAD_REQUEST, 'Unsupported WebSocket protocol version: %s' % ver)
        if self.method != 'GET':
            return self.simple_response(http_client.BAD_REQUEST, 'Invalid WebSocket method: %s' % self.method)

        response = HANDSHAKE_STR % as_base64_unicode(sha1((key + GUID_STR).encode('utf-8')).digest())
        self.optimize_for_sending_packet()
        self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.set_state(WRITE, self.upgrade_connection_to_ws, ReadOnlyFileBuffer(response.encode('ascii')), inheaders)

    def upgrade_connection_to_ws(self, buf, inheaders, event):
        if self.write(buf):
            if self.websocket_handler is None:
                self.websocket_handler = DummyHandler()
            self.read_frame, self.current_recv_opcode = ReadFrame(), None
            self.in_websocket_mode = True
            try:
                self.websocket_handler.handle_websocket_upgrade(self.websocket_connection_id, weakref.ref(self), inheaders)
            except Exception as err:
                self.log.exception('Error in WebSockets upgrade handler:')
                self.websocket_close(UNEXPECTED_ERROR, 'Unexpected error in handler: %r' % as_unicode(err))
            self.handle_event = self.ws_duplex
            self.set_ws_state()
            self.end_send_optimization()

    def set_ws_state(self):
        if self.ws_close_sent or self.ws_close_received:
            if self.ws_close_sent:
                self.ready = False
            else:
                self.wait_for = WRITE
            return

        if self.send_buf is not None or self.sending is not None:
            self.wait_for = RDWR
        else:
            try:
                self.sending = self.sendq.get_nowait()
            except Empty:
                with self.cf_lock:
                    if self.control_frames:
                        self.wait_for = RDWR
                    else:
                        self.wait_for = READ
            else:
                self.wait_for = RDWR

        if self.stop_reading:
            if self.wait_for is READ:
                self.ready = False
            elif self.wait_for is RDWR:
                self.wait_for = WRITE

    def ws_duplex(self, event):
        if event is READ:
            self.ws_read()
        elif event is WRITE:
            self.ws_write()
        self.set_ws_state()

    def ws_read(self):
        if not self.stop_reading:
            self.read_frame(self)

    def ws_data_received(self, data, opcode, frame_starting, frame_finished, is_final_frame_of_message):
        if opcode in CONTROL_CODES:
            return self.ws_control_frame(opcode, data)

        message_starting = self.current_recv_opcode is None
        if message_starting:
            if opcode == CONTINUATION:
                self.log.error('Client sent continuation frame with no message to continue')
                self.websocket_close(PROTOCOL_ERROR, 'Continuation frame without any message to continue')
                return
            self.current_recv_opcode = opcode
        elif frame_starting and opcode != CONTINUATION:
            self.log.error('Client sent continuation frame with non-zero opcode')
            self.websocket_close(PROTOCOL_ERROR, 'Continuation frame with non-zero opcode')
            return
        message_finished = frame_finished and is_final_frame_of_message
        if self.current_recv_opcode == TEXT:
            if message_starting:
                self.frag_decoder.reset()
            empty_data = len(data) == 0
            try:
                data = self.frag_decoder(data)
            except ValueError:
                self.frag_decoder.reset()
                self.log.error('Client sent undecodeable UTF-8')
                return self.websocket_close(INCONSISTENT_DATA, 'Not valid UTF-8')
            if message_finished:
                if (not data and not empty_data) or self.frag_decoder.state:
                    self.frag_decoder.reset()
                    self.log.error('Client sent undecodeable UTF-8')
                    return self.websocket_close(INCONSISTENT_DATA, 'Not valid UTF-8')
        if message_finished:
            self.current_recv_opcode = None
            self.frag_decoder.reset()
        try:
            self.handle_websocket_data(data, message_starting, message_finished)
        except Exception as err:
            self.log.exception('Error in WebSockets data handler:')
            self.websocket_close(UNEXPECTED_ERROR, 'Unexpected error in handler: %r' % as_unicode(err))

    def ws_control_frame(self, opcode, data):
        if opcode in (PING, CLOSE):
            rcode = PONG if opcode == PING else CLOSE
            if opcode == CLOSE:
                self.ws_close_received = True
                self.stop_reading = True
                if data:
                    try:
                        close_code = unpack_from(b'!H', data)[0]
                    except struct_error:
                        data = pack(b'!H', PROTOCOL_ERROR) + b'close frame data must be at least two bytes'
                    else:
                        try:
                            utf8_decode(data[2:])
                        except ValueError:
                            data = pack(b'!H', PROTOCOL_ERROR) + b'close frame data must be valid UTF-8'
                        else:
                            if close_code < 1000 or close_code in RESERVED_CLOSE_CODES or (1011 < close_code < 3000):
                                data = pack(b'!H', PROTOCOL_ERROR) + b'close code reserved'
                else:
                    close_code = NORMAL_CLOSE
                    data = pack(b'!H', close_code)
            f = ReadOnlyFileBuffer(create_frame(1, rcode, data))
            f.is_close_frame = opcode == CLOSE
            with self.cf_lock:
                self.control_frames.append(f)
        elif opcode == PONG:
            try:
                self.websocket_handler.handle_websocket_pong(self.websocket_connection_id, data)
            except Exception:
                self.log.exception('Error in PONG handler:')
        self.set_ws_state()

    def websocket_close(self, code=NORMAL_CLOSE, reason=b''):
        if isinstance(reason, str):
            reason = reason.encode('utf-8')
        self.stop_reading = True
        reason = reason[:123]
        if code is None and not reason:
            f = ReadOnlyFileBuffer(create_frame(1, CLOSE, b''))
        else:
            f = ReadOnlyFileBuffer(create_frame(1, CLOSE, pack(b'!H', code) + reason))
        f.is_close_frame = True
        with self.cf_lock:
            self.control_frames.append(f)
        self.set_ws_state()

    def ws_write(self):
        if self.ws_close_sent:
            return
        if self.send_buf is not None:
            if self.write(self.send_buf):
                self.end_send_optimization()
                if getattr(self.send_buf, 'is_close_frame', False):
                    self.ws_close_sent = True
                self.send_buf = None
        else:
            with self.cf_lock:
                try:
                    self.send_buf = self.control_frames.popleft()
                except IndexError:
                    if self.sending is not None:
                        self.send_buf = self.sending.create_frame()
                        if self.send_buf is None:
                            self.sending = None
            if self.send_buf is not None:
                self.optimize_for_sending_packet()

    def close(self):
        if self.in_websocket_mode:
            try:
                self.websocket_handler.handle_websocket_close(self.websocket_connection_id)
            except Exception:
                self.log.exception('Error in WebSocket close handler')
            # Try to write a close frame, just once
            try:
                if self.send_buf is None and not self.ws_close_sent:
                    self.websocket_close(SHUTTING_DOWN, 'Shutting down')
                    with self.cf_lock:
                        self.write(self.control_frames.pop())
            except Exception:
                pass
            Connection.close(self)
        else:
            HTTPConnection.close(self)
    # }}}

    def send_websocket_message(self, buf, wakeup=True):
        ''' Send a complete message. This class will take care of splitting it
        into appropriate frames automatically. `buf` must be a file like object. '''
        self.sendq.put(MessageWriter(buf))
        self.wait_for = RDWR
        if wakeup:
            self.wakeup()

    def send_websocket_frame(self, data, is_first=True, is_last=True):
        ''' Useful for streaming handlers that want to break up messages into
        frames themselves. Note that these frames will be interleaved with
        control frames, so they should not be too large. '''
        opcode = (TEXT if isinstance(data, str) else BINARY) if is_first else CONTINUATION
        fin = 1 if is_last else 0
        frame = create_frame(fin, opcode, data)
        with self.cf_lock:
            self.control_frames.append(ReadOnlyFileBuffer(frame))

    def send_websocket_ping(self, data=b''):
        ''' Send a PING to the remote client, it should reply with a PONG which
        will be sent to the handle_websocket_pong callback in your handler. '''
        if isinstance(data, str):
            data = data.encode('utf-8')
        frame = create_frame(True, PING, data)
        with self.cf_lock:
            self.control_frames.append(ReadOnlyFileBuffer(frame))

    def handle_websocket_data(self, data, message_starting, message_finished):
        ''' Called when some data is received from the remote client. In
        general the data may not constitute a complete "message", use the
        message_starting and message_finished flags to re-assemble it into a
        complete message in the handler. Note that for binary data, data is a
        mutable object. If you intend to keep it around after this method
        returns, create a bytestring from it, using tobytes(). '''
        self.websocket_handler.handle_websocket_data(self.websocket_connection_id, data, message_starting, message_finished)


class DummyHandler:

    def handle_websocket_upgrade(self, connection_id, connection_ref, inheaders):
        conn = connection_ref()
        conn.websocket_close(NORMAL_CLOSE, 'No WebSocket handler available')

    def handle_websocket_data(self, connection_id, data, message_starting, message_finished):
        pass

    def handle_websocket_pong(self, connection_id, data):
        pass

    def handle_websocket_close(self, connection_id):
        pass

# Testing {{{

# Run this file with calibre-debug and use wstest to run the Autobahn test
# suite


class EchoHandler:

    def __init__(self, *args, **kwargs):
        self.ws_connections = {}

    def conn(self, cid):
        ans = self.ws_connections.get(cid)
        if ans is not None:
            ans = ans()
        return ans

    def handle_websocket_upgrade(self, connection_id, connection_ref, inheaders):
        self.ws_connections[connection_id] = connection_ref

    def handle_websocket_data(self, connection_id, data, message_starting, message_finished):
        self.conn(connection_id).send_websocket_frame(data, message_starting, message_finished)

    def handle_websocket_pong(self, connection_id, data):
        pass

    def handle_websocket_close(self, connection_id):
        self.ws_connections.pop(connection_id, None)


def run_echo_server():
    s = ServerLoop(create_http_handler(websocket_handler=EchoHandler()))
    with HandleInterrupt(s.wakeup):
        s.serve_forever()


if __name__ == '__main__':
    # import cProfile
    # cProfile.runctx('r()', {'r':run_echo_server}, {}, filename='stats.profile')
    run_echo_server()
# }}}
