truf-server/tests/test_gcp_rsa_adc_hardening.py

import base64
import builtins
import json
from pathlib import Path
import sys
import unittest
from unittest import mock

import pandas as pd


ROOT = Path(__file__).resolve().parents[1]
APP_DIR = ROOT / 'app'
sys.path.insert(0, str(APP_DIR))

import dashboard
import scanner_db
from keycheckers.gcp import gcpKeycheck


def der_tlv(tag, value):
    length = len(value)
    if length < 0x80:
        encoded_length = bytes([length])
    else:
        length_bytes = length.to_bytes((length.bit_length() + 7) // 8, 'big')
        encoded_length = bytes([0x80 | len(length_bytes)]) + length_bytes
    return bytes([tag]) + encoded_length + value


def der_integer(value):
    value = value or b'\x00'
    if value[0] & 0x80:
        value = b'\x00' + value
    return der_tlv(0x02, value)


def rsa_sequence(bits=2048, modulus=None, private_exponent=b'\x01'):
    modulus = modulus or (b'\x80' + b'\x01' * (bits // 8 - 1))
    values = [
        der_integer(b'\x00'),
        der_integer(modulus),
        der_integer(b'\x01\x00\x01'),
        der_integer(private_exponent),
    ]
    values.extend(der_integer(b'\x03') for _ in range(5))
    return der_tlv(0x30, b''.join(values))


def private_key_pem(bits=2048, pkcs8=False, modulus=None, private_exponent=b'\x01'):
    rsa_der = rsa_sequence(bits, modulus, private_exponent)
    if pkcs8:
        algorithm = der_tlv(
            0x30,
            der_tlv(0x06, gcpKeycheck.RSA_ENCRYPTION_OID) + der_tlv(0x05, b''),
        )
        der = der_tlv(0x30, der_integer(b'\x00') + algorithm + der_tlv(0x04, rsa_der))
        label = 'PRIVATE KEY'
    else:
        der = rsa_der
        label = 'RSA PRIVATE KEY'
    body = base64.b64encode(der).decode('ascii')
    return f'-----BEGIN {label}-----\n{body}\n-----END {label}-----\n'


def credentials(pem):
    return {
        'client_email': 'fixture@example.invalid',
        'private_key_id': 'fixture-key-id',
        'private_key': pem,
        'project_id': 'fixture-project',
        'token_uri': gcpKeycheck.GOOGLE_TOKEN_URL,
    }


class GCPRSABoundsTests(unittest.TestCase):
    def assert_rejected_without_cpu_or_network(self, pem):
        with mock.patch.object(builtins, 'pow') as rsa_pow, \
                mock.patch.object(gcpKeycheck.requests, 'post') as post:
            result = gcpKeycheck.check_service_account(credentials(pem), None, 5)
        self.assertEqual(result['status'], 'DEAD')
        self.assertEqual(result['classification'], 'invalid_private_key')
        rsa_pow.assert_not_called()
        post.assert_not_called()

    def test_oversized_modulus_exponent_and_der_are_terminal_before_work(self):
        oversized_integer = b'\x01' * (gcpKeycheck.MAX_RSA_INTEGER_BYTES + 1)
        cases = {
            'modulus': private_key_pem(modulus=oversized_integer),
            'private_exponent': private_key_pem(private_exponent=oversized_integer),
            'der': (
                '-----BEGIN PRIVATE KEY-----\n'
                + base64.b64encode(b'\x00' * (gcpKeycheck.MAX_DER_BYTES + 1)).decode('ascii')
                + '\n-----END PRIVATE KEY-----\n'
            ),
        }
        for name, pem in cases.items():
            with self.subTest(name=name):
                self.assert_rejected_without_cpu_or_network(pem)

    def test_malformed_der_lengths_and_trailing_data_are_rejected(self):
        prefix = der_tlv(0x02, b'\x01' * 200)
        malformed = (
            der_tlv(0x30, prefix + b'\x30\x80'),
            der_tlv(0x30, prefix + b'\x30\x82\x01'),
            der_tlv(0x30, prefix + b'\x30\x83\x00\x00\x01'),
            der_tlv(0x30, prefix + b'\x30\x82\x40\x01'),
            rsa_sequence() + b'\x00',
        )
        for der in malformed:
            label = 'RSA PRIVATE KEY'
            pem = (
                f'-----BEGIN {label}-----\n'
                + base64.b64encode(der).decode('ascii')
                + f'\n-----END {label}-----\n'
            )
            with self.subTest(der=der[:8]):
                self.assert_rejected_without_cpu_or_network(pem)

    def test_normal_pkcs1_2048_and_pkcs8_4096_reach_token_request(self):
        response = mock.Mock(status_code=400)
        response.json.return_value = {'error': 'invalid_grant'}
        response.text = 'invalid_grant'
        for bits, pkcs8 in ((2048, False), (4096, True)):
            with self.subTest(bits=bits, pkcs8=pkcs8), \
                    mock.patch.object(builtins, 'pow', return_value=1) as rsa_pow, \
                    mock.patch.object(gcpKeycheck.requests, 'post', return_value=response) as post:
                result = gcpKeycheck.check_service_account(
                    credentials(private_key_pem(bits=bits, pkcs8=pkcs8)), None, 5
                )
            self.assertEqual(result['status'], 'DEAD')
            rsa_pow.assert_called_once()
            post.assert_called_once()


class ADCResourceRedactionTests(unittest.TestCase):
    def test_enrichment_keeps_refresh_token_out_of_resource(self):
        sentinel = 'ADC-REFRESH-TOKEN-SENTINEL'
        finding = {
            'DetectorName': 'GCPApplicationDefaultCredentials',
            'Raw': sentinel,
            'ScannerContext': {'nearby': json.dumps({
                'type': 'authorized_user',
                'client_id': 'public-client-id',
                'client_secret': 'client-secret',
                'refresh_token': sentinel,
                'quota_project_id': 'quota-project',
            })},
        }
        enriched = scanner_db.enrich_finding(finding)
        self.assertEqual(enriched['project_id'], 'quota-project')
        self.assertEqual(enriched['resource'], '')
        self.assertNotIn(sentinel, json.dumps(enriched))

        finding['PostmanContext'] = {
            'credential_kind': 'application_default_credentials',
            'json_path': sentinel,
        }
        self.assertEqual(scanner_db.enrich_finding(finding)['resource'], '')

    def test_render_redacts_only_adc_resource_rows(self):
        sentinel = 'HISTORICAL-ADC-REFRESH-TOKEN'
        frame = pd.DataFrame([
            {
                'detector_name': 'GCPApplicationDefaultCredentials',
                'credential_kind': '',
                'resource': sentinel,
            },
            {
                'detector_name': 'ImportedDetector',
                'credential_kind': 'application_default_credentials',
                'resource': sentinel + '-IMPORTED',
            },
            {
                'detector_name': 'AzureOpenAI',
                'credential_kind': 'azure_openai_key',
                'resource': 'azure-resource-name',
            },
            {
                'detector_name': 'GCP',
                'credential_kind': 'service_account',
                'resource': 'nonsecret-private-key-id',
            },
        ])
        with mock.patch.object(dashboard.st, 'dataframe') as render:
            dashboard.display_df(frame)
        rendered = render.call_args.args[0]
        rendered_text = rendered.to_string()
        self.assertNotIn(sentinel, rendered_text)
        self.assertIn('azure-resource-name', rendered_text)
        self.assertIn('nonsecret-private-key-id', rendered_text)


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