AO2XP/AOsocket.py

import socket
import platform

import websocket

from game_version import *


class AOtcpSocket(object):
    def __init__(self):
        self.sock = socket.socket()
        self.isWS = False

    def connect(self, ip, port):
        try:
            self.sock.connect((ip, port))
        except:
            return False

        self.sock.settimeout(0.1)
        return True

    def recv(self):
        tempdata = ""
        gotIt = False
        while not gotIt:
            try:
                contents = self.sock.recv(16384)
            except (socket.timeout, socket.error) as err:
                error = err.args[0]
                if error == "timed out" or error == 10035 or err.errno == 11:
                    return -2, []
                else: # connection lost
                    return -1, []

            if not contents.endswith("%"):
                tempdata += contents
                continue
            else:
                gotIt = True
                if tempdata:
                    contents = tempdata + contents
                    tempdata = ""

        totals = contents.split('%')
        del totals[-1]
        for i in range(len(totals)):
            totals[i] = totals[i].split("#")
            del totals[i][-1]

        return 0, totals

    def send(self, data):
        return self.sock.send(data)

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

class AOwebSocket(object):
    def __init__(self):
        self.sock = websocket.WebSocket()
        self.isWS = True
        self.header = {
            "User-Agent": "AO2XP %s, Python %s, %s %s %s" % (GAME_VERSION, platform.python_version(), platform.system(), platform.release(), platform.machine())
        }

    def connect(self, ip, port):
        try:
            self.sock.connect("ws://%s:%s" % (ip, port), header=self.header)
        except:
            return False

        self.sock.settimeout(0.1)
        return True

    def recv(self):
        try:
            contents = self.sock.recv()
        except websocket.WebSocketTimeoutException:
            return -2, []
        except websocket.WebSocketConnectionClosedException:
            return -1, []

        totals = contents.split('%')
        del totals[-1]
        for i in range(len(totals)):
            totals[i] = totals[i].split("#")
            del totals[i][-1]

        return 0, totals

    def send(self, data):
        return self.sock.send(str(data))

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