import unittest

from fgai.graylog_source import GraylogStreamSource


class _Client:
    def __init__(self):
        self.arguments = None

    def probe(self):
        return {"status": "connected"}

    def call_tool(self, _name, arguments):
        self.arguments = arguments
        return {"result": {"content": [{"type": "text", "text": '{"schema":[{"field":"client"},{"field":"server"},{"field":"result"}],"datarows":[["10.0.0.1","8.8.8.8","deny"],["10.0.0.2","8.8.8.8","accept"]]}' }]}}


class GraylogSourceTests(unittest.TestCase):
    def test_applies_custom_mapping_to_generic_stream_message(self):
        client = _Client()
        events, status = GraylogStreamSource(
            client, "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}'
        ).fetch()
        self.assertEqual(status["events_fetched"], 2)
        self.assertEqual(events[0].src_ip, "10.0.0.1")
        self.assertEqual(events[0].dst_ip, "8.8.8.8")
        self.assertEqual(events[0].action, "deny")
        self.assertEqual(client.arguments["streams"], ["vpn"])
        self.assertEqual(client.arguments["range_seconds"], 300)
        self.assertIn("client", client.arguments["fields"])
        self.assertEqual(client.arguments["offset"], 0)

    def test_requests_selected_profile_fields(self):
        client = _Client()
        GraylogStreamSource(client, "windows", profile_fields=("TargetUserName", "EventID")).fetch()
        self.assertIn("targetusername", client.arguments["fields"])
        self.assertIn("eventid", client.arguments["fields"])


if __name__ == "__main__":
    unittest.main()