from asysocks.unicomm.common.target import UniTarget, UniProto
from asysocks.unicomm.server import UniServer
from atds.network.packetizer import TDSPacketizer
from atds.relay.serverconnection import TDSRelayServerConnection
import traceback
import asyncio

async def log_callback(msg):
    print(msg)


class TDSRelayServerSettings:
    def __init__(self, gssapi_factory, log_callback=log_callback):
        self.gssapi_factory = gssapi_factory
        self.log_callback = log_callback
        self.sql_version_major = 15
        self.sql_version_minor = 0
        self.sql_version_build = 4123
        self.packet_size = 4096

    def create_new_gssapi(self, connection):
        gssapi = self.gssapi_factory()
        gssapi.set_connection_info(connection)
        return gssapi


class TDSRelayServer:
    def __init__(self, target, settings):
        self.target = target
        self.settings = settings
        self.server = None
        self.serving_task = None
        self.connections = {}
        self.connection_ctr = 0

    async def log(self, msg):
        if self.settings.log_callback is not None:
            await self.settings.log_callback(msg)

    async def __handle_connection(self, connection):
        try:
            self.connection_ctr += 1
            cid = self.connection_ctr
            await self.log('[TDSRELAY][%s][INF] Got new connection!' % cid)

            try:
                tds_connection = TDSRelayServerConnection(
                    self.settings, connection, connection_id=cid
                )
                self.connections[cid] = tds_connection
                connection_task = await tds_connection.run()
                await connection_task
                await self.log('[TDSRELAY][%s][INF] Connection end' % cid)
            except Exception as e:
                await self.log('[TDSRELAY][%s][ERR] %s' % (cid, e))

        except Exception as e:
            await self.log('[TDSRELAY][%s][ERR] Server error: %s' % (self.connection_ctr, e))
        finally:
            cid = self.connection_ctr
            if cid in self.connections:
                await self.connections[cid].stop()
                del self.connections[cid]

    async def run(self):
        self.server = UniServer(self.target, TDSPacketizer())
        return await self.server.serve_callback(self.__handle_connection)


async def test_relay_queue(rq):
    """Example relay queue consumer that uses relayed credentials to connect to a target MSSQL server"""
    try:
        from atds.connection import MSSQLConnection
        from atds.common.target import MSSQLTarget

        test_target = MSSQLTarget('10.10.10.2', port=1433)
        while True:
            item = await rq.get()
            print('[RELAY] Got relayed auth context: %s' % item)
            connection = MSSQLConnection(test_target, None, auth=item)
            ok, err = await connection.connect()
            if err is not None:
                print('[RELAY] MSSQL client login err: %s' % err)
                continue
            result = await connection.batch("SELECT @@version")
            print('[RELAY] Query result: %s' % result)

    except Exception as e:
        traceback.print_exc()


async def amain():
    try:
        from asyauth.protocols.spnego.relay.native import spnegorelay_ntlm_factory
        from asyauth.protocols.ntlm.relay.native import NTLMRelaySettings, ntlmrelay_factory

        auth_relay_queue = asyncio.Queue()
        x = asyncio.create_task(test_relay_queue(auth_relay_queue))
        target = UniTarget('0.0.0.0', 1433, UniProto.SERVER_TCP)

        settings = TDSRelayServerSettings(
            lambda: spnegorelay_ntlm_factory(auth_relay_queue, lambda: ntlmrelay_factory())
        )

        server = TDSRelayServer(target, settings)
        server_task = await server.run()
        await server_task
    except Exception as e:
        traceback.print_exc()

if __name__ == '__main__':
    asyncio.run(amain())

