# -*- coding: utf-8 -*-
# Copyright: (c) 2019, Jordan Borean (@jborean93) <jborean93@gmail.com>
# MIT License (see LICENSE or https://opensource.org/licenses/MIT)

import errno
import logging
import select
import socket
import struct
import threading
import timeit

from collections import (
    OrderedDict,
)

from smbprotocol.structure import (
    BytesField,
    IntField,
    Structure,
)

log = logging.getLogger(__name__)


class DirectTCPPacket(Structure):
    """
    [MS-SMB2] v53.0 2017-09-15

    2.1 Transport
    The Directory TCP transport packet header MUST have the following
    structure.
    """

    def __init__(self):
        self.fields = OrderedDict([
            ('stream_protocol_length', IntField(
                size=4,
                little_endian=False,
                default=lambda s: len(s['smb2_message']),
            )),
            ('smb2_message', BytesField(
                size=lambda s: s['stream_protocol_length'].get_value(),
            )),
        ])
        super(DirectTCPPacket, self).__init__()


class Tcp(object):

    MAX_SIZE = 16777215

    def __init__(self, server, port, timeout=None):
        self.server = server
        self.port = port
        self.timeout = timeout
        self.connected = False
        self._sock = None
        self._sock_lock = threading.Lock()
        self._close_lock = threading.Lock()

    def connect(self):
        with self._sock_lock:
            if not self.connected:
                log.info("Connecting to DirectTcp socket")
                try:
                    self._sock = socket.create_connection((self.server, self.port), timeout=self.timeout)
                except (OSError, socket.gaierror) as err:
                    raise ValueError("Failed to connect to '%s:%s': %s" % (self.server, self.port, str(err))) from err
                self._sock.settimeout(None)  # Make sure the socket is in blocking mode.
                self.connected = True

    def close(self):
        with self._sock_lock:
            if self.connected:
                log.info("Disconnecting DirectTcp socket")

                # Sending shutdown first will tell the recv thread (for both select and recv) that the socket has data
                # which returns b'' meaning it was closed.
                try:
                    self._sock.shutdown(socket.SHUT_RDWR)
                except socket.error as e:  # pragma: no cover
                    # Avoid collecting coverage here to avoid CI failing due to race condition differences
                    if e.errno != errno.ENOTCONN:
                        raise

                # This is even more special, we cannot close the socket if we are in the middle of a select or recv().
                # Doing so causes either a timeout (bad!) or bad fd descriptor (somewhat bad). By shutting down the
                # socket first then waiting until recv() has exited the critical select/recv section we can ensure we
                # gracefully handle client side socket closures both here and our recv thread without any extra waits
                # or exceptions.
                with self._close_lock:
                    self._sock.close()
                    self.connected = False

    def send(self, header):
        b_msg = header
        data_length = len(b_msg)
        if data_length > self.MAX_SIZE:
            raise ValueError("Data to be sent over Direct TCP size %d exceeds the max length allowed %d"
                             % (data_length, self.MAX_SIZE))

        tcp_packet = DirectTCPPacket()
        tcp_packet['smb2_message'] = b_msg

        data = tcp_packet.pack()

        with self._sock_lock:
            while data:
                sent = self._sock.send(data)
                data = data[sent:]

    def recv(self, timeout):
        # We don't need a lock for recv as the receiver is called from 1 thread.
        b_packet_size, timeout = self._recv(4, timeout)
        if not b_packet_size:
            return b''

        packet_size = struct.unpack(">L", b_packet_size)[0]

        return self._recv(packet_size, timeout)[0]

    def _recv(self, length, timeout):
        buffer = bytearray(length)
        offset = 0
        while offset < length:
            read_len = length - offset
            log.debug("Socket recv(%s) (total %s)" % (read_len, length))

            start_time = timeit.default_timer()

            with self._close_lock:
                if not self.connected:
                    # The socket was closed - need the no cover to avoid CI failing on race condition differences
                    return None, timeout    # pragma: no cover

                read = select.select([self._sock], [], [], max(timeout, 1))[0]
                timeout = timeout - (timeit.default_timer() - start_time)
                if not read:
                    log.debug("Socket recv(%s) timed out")
                    raise TimeoutError()

                try:
                    b_data = self._sock.recv(read_len)
                except socket.error as e:
                    # Windows will raise this error if the socket has been shutdown, Linux return returns an empty byte
                    # string so we just replicate that.
                    if e.errno not in [errno.ESHUTDOWN, errno.ECONNRESET]:
                        # Avoid collecting coverage here to avoid CI failing due to race condition differences
                        raise  # pragma: no cover
                    b_data = b''

            read_len = len(b_data)
            log.debug("Socket recv() returned %s bytes (total %s)" % (read_len, length))

            if read_len == 0:
                self.close()
                return None, timeout  # The socket has been shutdown

            buffer[offset:offset + read_len] = b_data
            offset += read_len

        return bytes(buffer), timeout
