rcon_client.py

import dataclasses
import enum
import getpass
import secrets
import struct


class AuthenticationFailure(ValueError):
    """Wrong password."""



class RconTypeEnum(enum.IntEnum):
    login    =  3
    command  =  2
    response =  0
    error    = -1


@dataclasses.dataclass
class RconRequest:
    type : RconTypeEnum
    data : bytes
    id   : bytes = None
    
    def __post_init__(self):
        for e in list(RconTypeEnum):
            if self.type == e:
                self.type = e
                break
        else:
            raise ValueError(f'{self.type} not valid RCON type')
        
        if self.id is None:
            self.id = secrets.randbelow(0x7fffffff)


    @classmethod
    def frombytes(cls, data):
        sz = len(data) - 4 - 4 - 1
        id, type, payload = struct.unpack(f'<ii{sz}Bx', data)
        return cls(type, payload, id)
    
    
    def tobytes(self):
        sz = len(self.data) + 4 +4 + 1 
        pkt = struct.pack('<iii{sz}B', sz, self.id, self.type, self.data)
        return pkt + '\0'
   
    def __str__(self):
        return self.data.decode('latin1')

    


class RconProtocol(asyncio.Protocol):
    def __init__(self, password, event_conlost=None):
        self.event_conlost = event_conlost
        self.password = password.encode()
    
    def login(self):
        login = RconRequest(self.password, type=RconTypeEnum.login)
        login = login.tobytes()
        self.transport.send(login)
    
    def connection_made(self, transport):
        self.transport = transport
        self.login()
    
    def data_received(self, data):
        r = RconRequest.frombytes(data)
        # handle login response
        if r.type == RconTypeEnum.login:
            if r.id == -1:
                raise AuthenticationFailure('\n\nWrong password')
            else:
                print('Access granted.')
        else:
            print(r)
                
    
    def eof_received(self):
        self.transport.close()
        if self.event_conlost is not None:
            self.event_conlost.set()



class RconClient:
    def __init__(self, password, *, host='localhost', port=25576):
        self.password = password
        self.host = (host, port)
        self.loop = asyncio.get_running_loop()
        self.server = None
        self.event_conlost = asyncio.Event()
        self.running_tasks = []
        
    
    async def serve_forever(self):
        self.reader, self.writer = await self.loop.create_connection(
            lambda: self.password, RconProtocol(self.event_conlost),
            *self.host
        )
        asyncio.create_task(self.get_input())
        asyncio.create_task(self.poll_server())
        
        
    
    async def get_input(self):
        
        


    
async def main():
    rcon = RconClient(getpass.getpass())
    await rcon.serve_forever()
    
    
    
if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass