kube-hunter/tests/discovery/test_azure.py

import json
import time
import requests_mock

from kube_hunter.core.events.event_handler import handler
from kube_hunter.modules.discovery.hosts import RunningAsPodEvent
from kube_hunter.modules.discovery.cloud.azure import (
    AzureInstanceMetadataServiceDiscovery,
    AzureMetadataApiExposed,
    AzureSubnetsDiscovery,
)

event_counter = 0


def test_TestAzureMetadataApi():
    global event_counter

    f = AzureInstanceMetadataServiceDiscovery(RunningAsPodEvent())

    with requests_mock.Mocker() as m:
        m.get("http://169.254.169.254/metadata/versions", status_code=404)
        f.execute()

    # We expect 0 triggers.because versions returned 404
    time.sleep(0.01)
    assert event_counter == 0
    event_counter = 0

    with requests_mock.Mocker() as m:
        m.get("http://169.254.169.254/metadata/versions", text=AzureApiResponses.make_versions_response())
        m.get(
            "http://169.254.169.254/metadata/instance?api-version=2017-08-01",
            text=AzureApiResponses.make_instance_response([("192.168.1.0", "24")]),
        )
        f.execute()

    # Expect 1 trigger
    time.sleep(0.01)
    assert event_counter == 1
    event_counter = 0

    # Test subnet extraction:
    versions_info = {"2017-08-01": AzureApiResponses.make_instance_response([("192.168.0.0", "24")], raw=False)}
    asd = AzureSubnetsDiscovery(AzureMetadataApiExposed(versions_info))
    assert asd.extract_azure_subnet() == "192.168.0.0/24"


class AzureApiResponses:
    @staticmethod
    def make_instance_response(subnets, raw=True):
        response = {
            "network": {
                "interface": [
                    {"ipv4": {"subnet": [{"address": address, "prefix": prefix} for address, prefix in subnets]}}
                ]
            }
        }

        if raw:
            response = json.dumps(response)
        return response

    @staticmethod
    def make_versions_response():
        return json.dumps(
            {
                "apiVersions": [
                    "2017-08-01",
                ]
            }
        )


@handler.subscribe(AzureMetadataApiExposed)
class TestAzureMetadataApiExposed:
    def __init__(self, event):
        global event_counter
        event_counter += 1