wiregui/tests/test_services.py at 0dd25f1336218862bd289bb0e40caa513f103eaa · bartei/wiregui · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
"""Tests for services — WireGuard event error handling and rule events."""

from unittest.mock import AsyncMock, patch

from wiregui.models.device import Device
from wiregui.models.rule import Rule
from wiregui.services.events import on_device_created, on_device_deleted, on_device_updated, on_rule_created, _device_allowed_ips


def _make_device(**kwargs) -> Device:
    defaults = dict(
        name="test",
        public_key="pk-test",
        preshared_key="psk-test",
        ipv4="10.3.2.5",
        ipv6="fd00::3:2:5",
        user_id="00000000-0000-0000-0000-000000000000",
    )
    defaults.update(kwargs)
    return Device(**defaults)


# --- _device_allowed_ips tests ---


def test_device_allowed_ips_basic():
    """Test _device_allowed_ips returns tunnel IPs with /32 and /128."""
    device = _make_device()
    ips = _device_allowed_ips(device)
    assert ips == ["10.3.2.5/32", "fd00::3:2:5/128"]


def test_device_allowed_ips_with_relay_subnets():
    """Test _device_allowed_ips includes relay subnets."""
    device = _make_device(allowed_subnets=["192.168.1.0/24", "10.20.0.0/16"])
    ips = _device_allowed_ips(device)
    assert ips == ["10.3.2.5/32", "fd00::3:2:5/128", "192.168.1.0/24", "10.20.0.0/16"]


def test_device_allowed_ips_ipv4_only():
    """Test _device_allowed_ips with only IPv4."""
    device = _make_device(ipv6=None)
    ips = _device_allowed_ips(device)
    assert ips == ["10.3.2.5/32"]


def test_device_allowed_ips_ipv6_only():
    """Test _device_allowed_ips with only IPv6."""
    device = _make_device(ipv4=None)
    ips = _device_allowed_ips(device)
    assert ips == ["fd00::3:2:5/128"]


def test_device_allowed_ips_relay_only():
    """Test _device_allowed_ips with only relay subnets (no tunnel IPs)."""
    device = _make_device(ipv4=None, ipv6=None, allowed_subnets=["192.168.1.0/24"])
    ips = _device_allowed_ips(device)
    assert ips == ["192.168.1.0/24"]


def test_device_allowed_ips_empty():
    """Test _device_allowed_ips with no IPs or subnets."""
    device = _make_device(ipv4=None, ipv6=None, allowed_subnets=[])
    ips = _device_allowed_ips(device)
    assert ips == []


# --- Events (WG disabled) ---


@patch("wiregui.services.events.get_settings")
@patch("wiregui.services.events.wireguard")
async def test_events_skip_when_wg_disabled(mock_wg, mock_settings):
    mock_settings.return_value.wg_enabled = False
    mock_wg.add_peer = AsyncMock()
    mock_wg.remove_peer = AsyncMock()

    device = _make_device()
    await on_device_created(device)
    await on_device_deleted(device)
    await on_device_updated(device)

    mock_wg.add_peer.assert_not_awaited()
    mock_wg.remove_peer.assert_not_awaited()


# --- Events (WG error handling) ---


@patch("wiregui.services.events.get_settings")
@patch("wiregui.services.events.firewall")
@patch("wiregui.services.events.wireguard")
async def test_on_device_created_handles_wg_error(mock_wg, mock_fw, mock_settings):
    mock_settings.return_value.wg_enabled = True
    mock_wg.add_peer = AsyncMock(side_effect=RuntimeError("wg failed"))
    mock_fw.add_device_jump_rule = AsyncMock()

    device = _make_device()
    # Should not raise — error is logged
    await on_device_created(device)


@patch("wiregui.services.events.get_settings")
@patch("wiregui.services.events.firewall")
@patch("wiregui.services.events.wireguard")
async def test_on_device_created_with_relay_subnets(mock_wg, mock_fw, mock_settings):
    """Test that device creation with relay subnets passes correct allowed_ips to WireGuard, adds routes, and configures firewall."""
    mock_settings.return_value.wg_enabled = True
    mock_wg.add_peer = AsyncMock()
    mock_wg.add_routes = AsyncMock()
    mock_fw.add_user_chain = AsyncMock()
    mock_fw.add_device_jump_rule = AsyncMock()

    device = _make_device(allowed_subnets=["192.168.1.0/24", "10.20.0.0/16"])
    await on_device_created(device)

    # Verify WireGuard peer was added with tunnel IPs + relay subnets
    mock_wg.add_peer.assert_awaited_once_with(
        public_key="pk-test",
        allowed_ips=["10.3.2.5/32", "fd00::3:2:5/128", "192.168.1.0/24", "10.20.0.0/16"],
        preshared_key="psk-test",
    )

    # Verify routes were added for relay subnets
    mock_wg.add_routes.assert_awaited_once_with(["192.168.1.0/24", "10.20.0.0/16"])

    # Verify firewall jump rule was added with relay subnets
    mock_fw.add_device_jump_rule.assert_awaited_once_with(
        "00000000-0000-0000-0000-000000000000",
        "10.3.2.5",
        "fd00::3:2:5",
        ["192.168.1.0/24", "10.20.0.0/16"],
    )


# --- Rule events ---


@patch("wiregui.services.events.get_settings")
@patch("wiregui.services.events.firewall")
async def test_on_rule_created_calls_apply_rule(mock_fw, mock_settings):
    mock_settings.return_value.wg_enabled = True
    mock_fw.apply_rule = AsyncMock()

    rule = Rule(
        action="accept",
        destination="10.0.0.0/8",
        port_type="tcp",
        port_range="80",
        user_id="00000000-0000-0000-0000-000000000000",
    )
    await on_rule_created(rule)

    mock_fw.apply_rule.assert_awaited_once_with(
        "00000000-0000-0000-0000-000000000000", "10.0.0.0/8", "accept", "tcp", "80",
    )