netv/util_test.py at main · jvdillon/netv · 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
"""Tests for util.py."""

from __future__ import annotations

from typing import Any

import json
import pathlib
import urllib.error

import pytest

import util


def test_atomic_write_json_replaces_file_and_preserves_permissions(tmp_path: pathlib.Path):
    path = tmp_path / "settings.json"
    path.write_text('{"old": true}')
    path.chmod(0o640)

    util.atomic_write_json(path, {"new": True})

    assert json.loads(path.read_text()) == {"new": True}
    assert path.stat().st_mode & 0o777 == 0o640
    assert list(tmp_path.iterdir()) == [path]


def _fake_request(url: str) -> Any:
    """Create a minimal request object for testing."""

    class _Req:
        full_url = url
        headers: dict[str, str] = {}
        data = None
        origin_req_host = "original.com"

        def get_method(self) -> str:
            return "GET"

    return _Req()


class TestSafeRedirectHandler:
    def test_handler_allows_http(self):
        handler = util._SafeRedirectHandler()
        req = _fake_request("http://original.com")
        result = handler.redirect_request(
            req,
            fp=None,
            code=302,
            msg="Found",
            headers={},
            newurl="http://redirect.com/path",
        )
        assert result is not None

    def test_handler_allows_https(self):
        handler = util._SafeRedirectHandler()
        req = _fake_request("https://original.com")
        result = handler.redirect_request(
            req,
            fp=None,
            code=302,
            msg="Found",
            headers={},
            newurl="https://secure.com/path",
        )
        assert result is not None

    def test_handler_rejects_file_scheme(self):
        handler = util._SafeRedirectHandler()
        req = _fake_request("http://original.com")
        with pytest.raises(urllib.error.URLError, match="Unsafe redirect scheme"):
            handler.redirect_request(
                req,
                fp=None,
                code=302,
                msg="Found",
                headers={},
                newurl="file:///etc/passwd",
            )

    def test_handler_rejects_data_scheme(self):
        handler = util._SafeRedirectHandler()
        req = _fake_request("http://original.com")
        with pytest.raises(urllib.error.URLError, match="Unsafe redirect scheme"):
            handler.redirect_request(
                req,
                fp=None,
                code=302,
                msg="Found",
                headers={},
                newurl="data:text/html,<script>alert(1)</script>",
            )

    def test_handler_rejects_javascript_scheme(self):
        handler = util._SafeRedirectHandler()
        req = _fake_request("http://original.com")
        with pytest.raises(urllib.error.URLError, match="Unsafe redirect scheme"):
            handler.redirect_request(
                req,
                fp=None,
                code=302,
                msg="Found",
                headers={},
                newurl="javascript:alert(1)",
            )


def test_redact_url_credentials():
    assert (
        util.redact_url_credentials(
            "https://provider.example/live/remote-user/remote-pass/42.m3u8"
        )
        == "https://provider.example/live/***/***/42.m3u8"
    )
    assert (
        util.redact_url_credentials(
            "https://remote-user:remote-pass@provider.example/stream"
            "?username=remote-user&password=remote-pass&channel=42"
        )
        == "https://***:***@provider.example/stream"
        "?username=%2A%2A%2A&password=%2A%2A%2A&channel=42"
    )


if __name__ == "__main__":
    from testing import run_tests

    run_tests(__file__)