SMANAGE/gvsds/tests/test_attack_detect.py

"""模块：AttackDetectMiddleware 单元测试

覆盖 _classify 状态码分类 / ENABLED 开关 / 状态码过滤 /
exempt_paths 前缀豁免 / _log_attack 写入 AbnormalUserLog（mock）。
"""
import uuid
from unittest import mock

from django.test import RequestFactory, SimpleTestCase

from gvsds import middleware as mw


class ClassifyTestCase(SimpleTestCase):
    """_classify 状态码 → 攻击类型映射。"""

    def testClassify_401_AuthFailure(self):
        self.assertEqual(
            mw.AttackDetectMiddleware._classify(401),
            'auth_failure',
        )

    def testClassify_403_PermissionDenied(self):
        self.assertEqual(
            mw.AttackDetectMiddleware._classify(403),
            'permission_denied',
        )

    def testClassify_429_RateLimited(self):
        self.assertEqual(
            mw.AttackDetectMiddleware._classify(429),
            'rate_limited',
        )

    def testClassify_UnknownStatus_Suspicious(self):
        self.assertEqual(
            mw.AttackDetectMiddleware._classify(500),
            'suspicious',
        )

    def testClassify_400_Suspicious(self):
        self.assertEqual(
            mw.AttackDetectMiddleware._classify(400),
            'suspicious',
        )


class AttackDetectMiddlewareTestCase(SimpleTestCase):
    """AttackDetectMiddleware 核心流程。"""

    def setUp(self):
        self.factory = RequestFactory()

    def _BuildMiddleware(self, enabled=True, status_codes=None,
                         exempt_paths=None, max_body_len=1024):
        """构造中间件实例，绕过 app.ini 读取。"""
        mwf = mw.AttackDetectMiddleware.__new__(mw.AttackDetectMiddleware)
        mwf.get_response = lambda req: _FakeResponse(200)
        mwf.enabled = enabled
        mwf.status_codes = set(status_codes or [401, 403, 429])
        mwf.exempt_paths = tuple(exempt_paths or ['/API/Auth', '/health'])
        mwf.max_body_len = max_body_len
        #: 绕过 __init__ 时需补齐限频去重状态，否则 _log_attack 内 AttributeError 被吞
        mwf._recent_attacks = {}
        return mwf

    def testDisabled_PassesThrough(self):
        mwf = self._BuildMiddleware(enabled=False)
        req = self.factory.get('/data', REMOTE_ADDR='1.2.3.4')
        resp = mwf(req)
        self.assertEqual(resp.status_code, 200)

    def testNormalStatus_PassesThrough(self):
        mwf = self._BuildMiddleware()
        # 200 不在检测列表 → 不记录
        req = self.factory.get('/data', REMOTE_ADDR='1.2.3.4')
        resp = mwf(req)
        self.assertEqual(resp.status_code, 200)

    def testExemptPath_PassesThrough(self):
        mwf = self._BuildMiddleware()
        # 构造一个 401 响应触发检测，但路径在豁免列表
        mwf.get_response = lambda req: _FakeResponse(401)
        req = self.factory.get('/API/Auth/Login', REMOTE_ADDR='1.2.3.4')
        resp = mwf(req)
        # 豁免路径 → 不记录，直接返回原 401
        self.assertEqual(resp.status_code, 401)

    def testExemptPath_ExactMatch_PassesThrough(self):
        mwf = self._BuildMiddleware()
        mwf.get_response = lambda req: _FakeResponse(401)
        req = self.factory.get('/health', REMOTE_ADDR='1.2.3.4')
        resp = mwf(req)
        self.assertEqual(resp.status_code, 401)

    def testDetectedStatus_CallsLogAttack(self):
        mwf = self._BuildMiddleware()
        mwf.get_response = lambda req: _FakeResponse(401)
        req = self.factory.get('/data', REMOTE_ADDR='1.2.3.4')
        with mock.patch.object(mwf, '_log_attack') as mock_log:
            mwf(req)
            mock_log.assert_called_once()

    def testNonExemptPath_403_CallsLogAttack(self):
        mwf = self._BuildMiddleware()
        mwf.get_response = lambda req: _FakeResponse(403)
        req = self.factory.get('/secure/data', REMOTE_ADDR='1.2.3.4')
        with mock.patch.object(mwf, '_log_attack') as mock_log:
            mwf(req)
            mock_log.assert_called_once()

    def testCustomStatusCodes_500_NotTriggered(self):
        mwf = self._BuildMiddleware(status_codes=[401, 403])
        mwf.get_response = lambda req: _FakeResponse(500)
        req = self.factory.get('/data', REMOTE_ADDR='1.2.3.4')
        with mock.patch.object(mwf, '_log_attack') as mock_log:
            mwf(req)
            mock_log.assert_not_called()

    def testCustomStatusCodes_429_Triggered(self):
        mwf = self._BuildMiddleware(status_codes=[429])
        mwf.get_response = lambda req: _FakeResponse(429)
        req = self.factory.get('/data', REMOTE_ADDR='1.2.3.4')
        with mock.patch.object(mwf, '_log_attack') as mock_log:
            mwf(req)
            mock_log.assert_called_once()


class LogAttackTestCase(SimpleTestCase):
    """_log_attack 写入 AbnormalUserLog（mock 数据库）。"""

    def setUp(self):
        self.factory = RequestFactory()

    def _BuildMiddleware(self, max_body_len=1024):
        mwf = mw.AttackDetectMiddleware.__new__(mw.AttackDetectMiddleware)
        mwf.max_body_len = max_body_len
        #: 绕过 __init__ 时需补齐限频去重状态，否则 _log_attack 内 AttributeError 被吞
        mwf._recent_attacks = {}
        return mwf

    def testLogAttack_WritesToAbnormalUserLog(self):
        mwf = self._BuildMiddleware()
        req = self.factory.get(
            '/secure/data',
            REMOTE_ADDR='1.2.3.4',
            HTTP_USER_AGENT='BadBot/1.0',
            HTTP_X_REAL_IP='5.6.7.8',
        )
        req.user = None
        resp = _FakeResponse(401)

        with mock.patch('gvsds.apps.gerp.models.AbnormalUserLog') as MockLog:
            MockManager = MockLog.objects.using.return_value
            mwf._log_attack(req, resp)
            MockManager.create.assert_called_once()
            CallKwargs = MockManager.create.call_args.kwargs
            self.assertEqual(CallKwargs['IPAddress'], '1.2.3.4')
            self.assertEqual(CallKwargs['RealIP'], '5.6.7.8')
            self.assertEqual(CallKwargs['UserAgent'], 'BadBot/1.0')
            self.assertEqual(CallKwargs['RequestMethod'], 'GET')
            self.assertEqual(CallKwargs['RequestPath'], '/secure/data')
            self.assertEqual(CallKwargs['ResponseStatus'], 401)
            self.assertEqual(CallKwargs['AttemptType'], 'auth_failure')

    def testLogAttack_WithRequestUser_ExtractsUserUUID(self):
        mwf = self._BuildMiddleware()
        req = self.factory.get('/data', REMOTE_ADDR='1.2.3.4')
        UserUUID = uuid.uuid4().bytes
        FakeUser = type('FakeUser', (), {'UserUUID': UserUUID})()
        req.user = FakeUser
        resp = _FakeResponse(403)

        with mock.patch('gvsds.apps.gerp.models.AbnormalUserLog') as MockLog:
            MockManager = MockLog.objects.using.return_value
            mwf._log_attack(req, resp)
            CallKwargs = MockManager.create.call_args.kwargs
            self.assertEqual(CallKwargs['UserUUID'], UserUUID)
            self.assertEqual(CallKwargs['AttemptType'], 'permission_denied')

    def testLogAttack_PostWithBody_CapturesBodySummary(self):
        mwf = self._BuildMiddleware(max_body_len=64)
        req = self.factory.post(
            '/data',
            data={'secret': 'leaked'},
            REMOTE_ADDR='1.2.3.4',
        )
        req.user = None
        resp = _FakeResponse(429)

        with mock.patch('gvsds.apps.gerp.models.AbnormalUserLog') as MockLog:
            MockManager = MockLog.objects.using.return_value
            mwf._log_attack(req, resp)
            CallKwargs = MockManager.create.call_args.kwargs
            self.assertEqual(CallKwargs['AttemptType'], 'rate_limited')
            # 请求体摘要应非空
            self.assertTrue(CallKwargs['RequestBody'])

    def testLogAttack_Exception_SwallowsAndLogs(self):
        """_log_attack 内部异常不应影响响应返回。"""
        mwf = self._BuildMiddleware()
        req = self.factory.get('/data', REMOTE_ADDR='1.2.3.4')
        req.user = None
        resp = _FakeResponse(401)

        with mock.patch('gvsds.apps.gerp.models.AbnormalUserLog') as MockLog:
            MockLog.objects.using.side_effect = RuntimeError('db down')
            # 不应抛异常
            mwf._log_attack(req, resp)

    def testLogAttack_SafeHeaders_OnlyWhitelistedKeys(self):
        """请求头摘要仅保留安全相关头，不记录 Authorization 等。"""
        mwf = self._BuildMiddleware()
        req = self.factory.get(
            '/data',
            REMOTE_ADDR='1.2.3.4',
            HTTP_USER_AGENT='Bot',
            HTTP_HOST='example.com',
            HTTP_REFERER='http://evil.com',
            HTTP_AUTHORIZATION='Bearer secret-token',
        )
        req.user = None
        resp = _FakeResponse(401)

        with mock.patch('gvsds.apps.gerp.models.AbnormalUserLog') as MockLog:
            MockManager = MockLog.objects.using.return_value
            mwf._log_attack(req, resp)
            CallKwargs = MockManager.create.call_args.kwargs
            Headers = CallKwargs['RequestHeaders']
            # 安全相关头应被保留
            self.assertIn('HTTP_USER_AGENT', Headers)
            self.assertIn('HTTP_HOST', Headers)
            # Authorization 不应出现在摘要中
            self.assertNotIn('HTTP_AUTHORIZATION', Headers)


class _FakeResponse:
    """简易响应对象。"""
    def __init__(self, status_code):
        self.status_code = status_code