# minestat.py - A Minecraft server status checker
# Copyright (C) 2016-2021 Lloyd Dilley, Felix Ern (MindSolve)
# http://www.dilley.me/
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
import json
import socket
import struct
from time import perf_counter
from enum import Enum
from typing import Union

class ConnStatus(Enum):
  """
Contains possible connection states.

- `SUCCESS`: The specified SLP connection succeeded (Request & response parsing OK)
- `CONNFAIL`: The socket to the server could not be established. Server offline, wrong hostname or port?
- `TIMEOUT`: The connection timed out. (Server under too much load? Firewall rules OK?)
- `UNKNOWN`: The connection was established, but the server spoke an unknown/unsupported SLP protocol.
  """

  def __str__(self):
    return str(self.name)

  SUCCESS = 0
  """The specified SLP connection succeeded (Request & response parsing OK)"""

  CONNFAIL = -1
  """The socket to the server could not be established. (Server offline, wrong hostname or port?)"""

  TIMEOUT = -2
  """The connection timed out. (Server under too much load? Firewall rules OK?)"""

  UNKNOWN = -3
  """The connection was established, but the server spoke an unknown/unsupported SLP protocol."""

class SlpProtocols(Enum):
  """
Contains possible SLP (Server List Ping) protocols.

- `JSON`: The newest and currently supported SLP protocol.

  Uses (wrapped) JSON as payload. Complex query, see `json_query()` for the protocol implementation.

  *Available since Minecraft 1.7*
- `EXTENDED_LEGACY`: The previous SLP protocol

  Used by Minecraft 1.6, it is still supported by all newer server versions.
  Complex query needed, see implementation `extended_legacy_query()` for full protocol details.

  *Available since Minecraft 1.6*
- `LEGACY`: The legacy SLP protocol.

  Used by Minecraft 1.4 and 1.5, it is the first protocol to contain the server version number.
  Very simple protocol call (2 byte), simple response decoding.
  See `legacy_query()` for full implementation and protocol details.

  *Available since Minecraft 1.4*
- `BETA`: The first SLP protocol.

  Used by Minecraft Beta 1.8 till Release 1.3, it is the first SLP protocol.
  It contains very few details, no server version info, only MOTD, max- and online player counts.

  *Available since Minecraft Beta 1.8*
  """

  def __str__(self):
    return str(self.name)

  JSON = 3
  """
  The newest and currently supported SLP protocol.

  Uses (wrapped) JSON as payload. Complex query, see `json_query()` for the protocol implementation.

  *Available since Minecraft 1.7*
  """

  EXTENDED_LEGACY = 2
  """The previous SLP protocol

  Used by Minecraft 1.6, it is still supported by all newer server versions.
  Complex query needed, see implementation `extended_legacy_query()` for full protocol details.

  *Available since Minecraft 1.6*
  """

  LEGACY = 1
  """
  The legacy SLP protocol.

  Used by Minecraft 1.4 and 1.5, it is the first protocol to contain the server version number.
  Very simple protocol call (2 byte), simple response decoding.
  See `legacy_query()` for full implementation and protocol details.

  *Available since Minecraft 1.4*
  """

  BETA = 0
  """
  The first SLP protocol.

  Used by Minecraft Beta 1.8 till Release 1.3, it is the first SLP protocol.
  It contains very few details, no server version info, only MOTD, max- and online player counts.

  *Available since Minecraft Beta 1.8*
  """

class MineStat:
  VERSION = "2.1.2"             # MineStat version
  DEFAULT_TIMEOUT = 5           # default TCP timeout in seconds

  def __init__(self, address, port, timeout = DEFAULT_TIMEOUT, query_protocol: SlpProtocols = None):
    self.address = address
    self.port = port
    self.online = None           # online or offline?
    self.version = None          # server version
    self.motd = None             # message of the day
    self.current_players = None  # current number of players online
    self.max_players = None      # maximum player capacity
    self.latency = None          # ping time to server in milliseconds
    self.timeout = timeout       # socket timeout
    self.slp_protocol = None     # Server List Ping protocol

    # Future improvement: IPv4/IPv6, multiple addresses
    # If a host has multiple IP addresses or a IPv4 and a IPv6 address,
    # socket.connect choses the first IPv4 address returned by DNS.
    # If a mc server is not available over IPv4, this failes as "offline".
    # Or in some environments, the DNS returns the external and the internal
    # address, but from an internal client, only the internal address is reachable
    # See https://docs.python.org/3/library/socket.html#socket.getaddrinfo

    # If the user wants a specific protocol, use only that.
    if query_protocol:
      if query_protocol is SlpProtocols.BETA:
        self.beta_query()
      elif query_protocol is SlpProtocols.LEGACY:
        self.legacy_query()
      elif query_protocol is SlpProtocols.EXTENDED_LEGACY:
        self.extended_legacy_query()
      elif query_protocol is SlpProtocols.JSON:
        self.json_query()

      return

    # Note: The order here is unfortunately important.
    # Some older versions of MC don't accept packets for a few seconds
    # after receiving a not understood packet.
    # An example is MC 1.4: Nothing works directly after a json request.
    # A legacy query alone works fine.

    # Minecraft 1.4 & 1.5 (legacy SLP)
    result = self.legacy_query()

    # Minecraft Beta 1.8 to Release 1.3 (beta SLP)
    if result is not ConnStatus.CONNFAIL \
        and result is not ConnStatus.SUCCESS:
      result = self.beta_query()

    # Minecraft 1.6 (extended legacy SLP)
    if result is not ConnStatus.CONNFAIL:
      result = self.extended_legacy_query()

    # Minecraft 1.7+ (JSON SLP)
    if result is not ConnStatus.CONNFAIL:
      self.json_query()

  def json_query(self):
    """
    Method for querying a modern (MC Java >= 1.7) server with the SLP protocol.
    This protocol is based on encoded JSON, see the documentation at wiki.vg below
    for a full packet description.

    See https://wiki.vg/Server_List_Ping#Current
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(self.timeout)

    try:
      start_time = perf_counter()
      sock.connect((self.address, self.port))
      self.latency = round((perf_counter() - start_time) * 1000)
    except socket.timeout:
      return ConnStatus.TIMEOUT
    except OSError:
      return ConnStatus.CONNFAIL

    # Construct Handshake packet
    req_data = bytearray([0x00])
    # Add protocol version. If pinging to determine version, use `-1`
    req_data += bytearray([0xff, 0xff, 0xff, 0xff, 0x0f])
    # Add server address length
    req_data += self._pack_varint(len(self.address))
    # Server address. Encoded with UTF8
    req_data += bytearray(self.address, "utf8")
    # Server port
    req_data += struct.pack(">H", self.port)
    # Next packet state (1 for status, 2 for login)
    req_data += bytearray([0x01])

    # Prepend full packet length
    req_data = self._pack_varint(len(req_data)) + req_data

    # Now actually send the constructed client request
    sock.send(req_data)

    # Now send empty "Request" packet
    # varint len, 0x00
    sock.send(bytearray([0x01, 0x00]))

    # Do all the receiving in a try-catch, to reduce duplication of error handling
    try:
      # Receive answer: full packet length as varint
      packet_len = self._unpack_varint(sock)

      # Check if full packet length seems acceptable
      if packet_len < 3:
        return ConnStatus.UNKNOWN

      # Receive actual packet id
      packet_id = self._unpack_varint(sock)

      # If we receive a packet with id 0x19, something went wrong.
      # Usually the payload is JSON text, telling us what exactly.
      # We could stop here, and display something to the user, as this is not normal
      # behaviour, maybe a bug somewhere here.

      # Instead I am just going to check for the correct packet id: 0x00
      if packet_id != 0:
        return ConnStatus.UNKNOWN

      # Receive & unpack payload length
      content_len = self._unpack_varint(sock)

      # Receive full payload
      payload_raw = self._recv_exact(sock, content_len)

    except socket.timeout:
      return ConnStatus.TIMEOUT
    except (ConnectionResetError, ConnectionAbortedError):
      return ConnStatus.UNKNOWN
    except OSError:
      return ConnStatus.CONNFAIL
    finally:
      sock.close()

    # Set protocol version
    self.slp_protocol = SlpProtocols.JSON

    # Parse and save to object attributes
    return self.__parse_json_payload(payload_raw)

  def __parse_json_payload(self, payload_raw: Union[bytes, bytearray]) -> ConnStatus:
    """
    Helper method for parsing the modern JSON-based SLP protocol.
    In use for Minecraft Java >= 1.7, see `json_query()` above for details regarding the protocol.

    :param payload_raw: The raw SLP payload, without header and string lenght
    """
    try:
      payload_obj = json.loads(payload_raw.decode('utf8'))
    except json.JSONDecodeError:
      return ConnStatus.UNKNOWN

    # Now that we have the status object, set all fields
    self.version = payload_obj["version"]["name"]

    if type(payload_obj["description"]) is str:
        self.motd = payload_obj["description"]
    elif type(payload_obj["description"]) is dict and "text" in payload_obj["description"]:
        self.motd = payload_obj["description"]["text"]

    self.max_players = payload_obj["players"]["max"]
    self.current_players = payload_obj["players"]["online"]

    # If we got here, everything is in order.
    self.online = True
    return ConnStatus.SUCCESS

  def _unpack_varint(self, sock: socket.socket):
    """ Small helper method for unpacking an int from an varint (streamed from socket). """
    data = 0
    for i in range(5):
      ordinal = sock.recv(1)

      if len(ordinal) == 0:
        break

      byte = ord(ordinal)
      data |= (byte & 0x7F) << 7 * i

      if not byte & 0x80:
        break

    return data

  def _pack_varint(self, data):
    """ Small helper method for packing a varint from an int. """
    ordinal = b''

    while True:
      byte = data & 0x7F
      data >>= 7
      ordinal += struct.pack('B', byte | (0x80 if data > 0 else 0))

      if data == 0:
        break

    return ordinal

  def extended_legacy_query(self):
    """
    Minecraft 1.6 SLP query, extended legacy ping protocol.
    All modern servers are currently backwards compatible with this protocol.

    See https://wiki.vg/Server_List_Ping#1.6
    :return:
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(self.timeout)

    try:
      start_time = perf_counter()
      sock.connect((self.address, self.port))
      self.latency = round((perf_counter() - start_time) * 1000)
    except socket.timeout:
      return ConnStatus.TIMEOUT
    except OSError:
      return ConnStatus.CONNFAIL

    # Send 0xFE as packet identifier,
    # 0x01 as ping packet content
    # 0xFA as packet identifier for a plugin message
    # 0x00 0x0B as strlen of following string
    req_data = bytearray([0xFE, 0x01, 0xFA, 0x00, 0x0B])
    # the string 'MC|PingHost' as UTF-16BE encoded string
    req_data += bytearray("MC|PingHost", "utf-16-be")
    # 0xXX 0xXX byte count of rest of data, 7+len(serverhostname), as short
    req_data += struct.pack(">h", 7 + (len(self.address) * 2))
    # 0xXX [legacy] protocol version (before netty rewrite)
    # Used here: 74 (MC 1.6.2)
    req_data += bytearray([0x49])
    # strlen of serverhostname (big-endian short)
    req_data += struct.pack(">h", len(self.address))
    # the hostname of the server
    req_data += bytearray(self.address, "utf-16-be")
    # port of the server, as int (4 byte)
    req_data += struct.pack(">i", self.port)

    # Now send the contructed client requests
    sock.send(req_data)

    try:
      # Receive answer packet id (1 byte)
      packet_id = self._recv_exact(sock, 1)

      # Check packet id (should be "kick packet 0xFF")
      if packet_id[0] != 0xFF:
        return ConnStatus.UNKNOWN

      # Receive payload lengh (signed big-endian short; 2 byte)
      raw_payload_len = self._recv_exact(sock, 2)

      # Extract payload length
      # Might be empty, if the server keeps the connection open but doesn't send anything
      content_len = struct.unpack(">h", raw_payload_len)[0]

      # Check if payload length is acceptable
      if content_len < 3:
        return ConnStatus.UNKNOWN

      # Receive full payload and close socket
      payload_raw = self._recv_exact(sock, content_len * 2)

    except socket.timeout:
      return ConnStatus.TIMEOUT
    except (ConnectionResetError, ConnectionAbortedError, struct.error):
      return ConnStatus.UNKNOWN
    except OSError:
      return ConnStatus.CONNFAIL
    finally:
      sock.close()

    # Set protocol version
    self.slp_protocol = SlpProtocols.EXTENDED_LEGACY

    # Parse and save to object attributes
    return self.__parse_legacy_payload(payload_raw)

  def legacy_query(self):
    """
    Minecraft 1.4-1.5 SLP query, server response contains more info than beta SLP

    See https://wiki.vg/Server_List_Ping#1.4_to_1.5

    :return: ConnStatus
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(self.timeout)

    try:
      start_time = perf_counter()
      sock.connect((self.address, self.port))
      self.latency = round((perf_counter() - start_time) * 1000)
    except socket.timeout:
      return ConnStatus.TIMEOUT
    except OSError:
      return ConnStatus.CONNFAIL

    # Send 0xFE 0x01 as packet id
    sock.send(bytearray([0xFE, 0x01]))

    # Receive answer packet id (1 byte) and payload lengh (signed big-endian short; 2 byte)
    try:
      raw_header = self._recv_exact(sock, 3)
    except socket.timeout:
      return ConnStatus.TIMEOUT
    except (ConnectionAbortedError, ConnectionResetError):
      return ConnStatus.UNKNOWN
    except OSError:
      return ConnStatus.CONNFAIL

    # Extract payload length
    # Might be empty, if the server keeps the connection open but doesn't send anything
    try:
      content_len = struct.unpack(">xh", raw_header)[0]
    except struct.error:
      return ConnStatus.UNKNOWN

    # Receive full payload and close socket
    payload_raw = bytearray(self._recv_exact(sock, content_len * 2))
    sock.close()

    # Set protocol version
    self.slp_protocol = SlpProtocols.LEGACY

    # Parse and save to object attributes
    return self.__parse_legacy_payload(payload_raw)

  def __parse_legacy_payload(self, payload_raw: Union[bytearray, bytes]) -> ConnStatus:
    """
    Internal helper method for parsing the legacy SLP payload (legacy and extended legacy).

    :param payload_raw: The extracted legacy SLP payload as bytearray/bytes
    """
    # According to wiki.vg, beta, legacy and extended legacy use UTF-16BE as "payload" encoding
    payload_str = payload_raw.decode('utf-16-be')

    # This "payload" contains six fields delimited by a NUL character:
    # - a fixed prefix '§1'
    # - the protocol version
    # - the server version
    # - the MOTD
    # - the online player count
    # - the max player count
    payload_list = payload_str.split('\x00')

    # Check for count of string parts, expected is 6 for this protocol version
    if len(payload_list) != 6:
      return ConnStatus.UNKNOWN

    # - a fixed prefix '§1'
    # - the protocol version
    # - the server version
    self.version = payload_list[2]
    # - the MOTD
    self.motd = payload_list[3]
    # - the online player count
    self.current_players = int(payload_list[4])
    # - the max player count
    self.max_players = int(payload_list[5])

    # If we got here, everything is in order
    self.online = True
    return ConnStatus.SUCCESS

  def beta_query(self):
    """
    Minecraft Beta 1.8 to Release 1.3 SLP protocol
    See https://wiki.vg/Server_List_Ping#Beta_1.8_to_1.3

    :return: ConnStatus
    """

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(self.timeout)

    try:
      start_time = perf_counter()
      sock.connect((self.address, self.port))
      self.latency = round((perf_counter() - start_time) * 1000)
    except socket.timeout:
      return ConnStatus.TIMEOUT
    except OSError:
      return ConnStatus.CONNFAIL

    # Send 0xFE as packet id
    sock.send(bytearray([0xFE]))

    # Receive answer packet id (1 byte) and payload lengh (signed big-endian short; 2 byte)
    try:
      raw_header = self._recv_exact(sock, 3)
    except socket.timeout:
      return ConnStatus.TIMEOUT
    except (ConnectionResetError, ConnectionAbortedError):
      return ConnStatus.UNKNOWN
    except OSError:
      return ConnStatus.CONNFAIL

    # Extract payload length
    # Might be empty, if the server keeps the connection open but doesn't send anything
    try:
      content_len = struct.unpack(">xh", raw_header)[0]
    except struct.error:
      return ConnStatus.UNKNOWN

    # Receive full payload and close socket
    payload_raw = bytearray(self._recv_exact(sock, content_len * 2))
    sock.close()

    # Set protocol version
    self.slp_protocol = SlpProtocols.BETA

    # According to wiki.vg, beta, legacy and extended legacy use UTF-16BE as "payload" encoding
    payload_str = payload_raw.decode('utf-16-be')
    # This "payload" contains three values:
    # The MOTD, the max player count, and the online player count
    payload_list = payload_str.split('§')

    # Check for count of string parts, expected is 3 for this protocol version
    # Note: We could check here if the list has the len() one, as that is most probably an error message.
    # e.g. ['Protocol error']
    if len(payload_list) < 3:
      return ConnStatus.UNKNOWN

    # The last value is the max player count
    self.max_players = int(payload_list[-1])
    # The second(-to-last) value is the online player count
    self.current_players = int(payload_list[-2])
    # The first value it the server MOTD
    # This could contain '§' itself, thats the reason for the join here
    self.motd = "§".join(payload_list[:-2])

    # Set general version, as the protocol doesn't contain the server version
    self.version = ">=1.8b/1.3"

    # If we got here, everything is in order
    self.online = True

    return ConnStatus.SUCCESS

  @staticmethod
  def _recv_exact(sock: socket.socket, size: int) -> bytearray:
    """
    Helper function for receiving a specific amount of data. Works around the problems of `socket.recv`.
    Throws a ConnectionAbortedError if the connection was closed while waiting for data.

    :param sock: Open socket to receive data from
    :param size: Amount of bytes of data to receive
    :return: bytearray with the received data
    """
    data = bytearray()

    while len(data) < size:
      temp_data = bytearray(sock.recv(size - len(data)))

      # If the connection was closed, `sock.recv` returns an empty string
      if not temp_data:
        raise ConnectionAbortedError

      data += temp_data

    return data
