Repository navigation
Expand file tree
/
Copy pathtest_security.py
More file actions
187 lines (150 loc) · 6.7 KB
/
Copy pathtest_security.py
File metadata and controls
187 lines (150 loc) · 6.7 KB
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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
"""Tests for graphify/security.py - URL validation, safe fetch, path guards, label sanitisation."""
from __future__ import annotations
import json
import urllib.error
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from graphify.security import (
sanitize_label,
safe_fetch,
safe_fetch_text,
validate_graph_path,
validate_url,
_MAX_FETCH_BYTES,
_MAX_TEXT_BYTES,
)
# ---------------------------------------------------------------------------
# validate_url
# ---------------------------------------------------------------------------
def test_validate_url_accepts_http():
assert validate_url("http://example.com/page") == "http://example.com/page"
def test_validate_url_accepts_https():
assert validate_url("https://arxiv.org/abs/1706.03762") == "https://arxiv.org/abs/1706.03762"
def test_validate_url_rejects_file():
with pytest.raises(ValueError, match="file"):
validate_url("file:///etc/passwd")
def test_validate_url_rejects_ftp():
with pytest.raises(ValueError, match="ftp"):
validate_url("ftp://files.example.com/data.zip")
def test_validate_url_rejects_data():
with pytest.raises(ValueError, match="data"):
validate_url("data:text/html,<script>alert(1)</script>")
def test_validate_url_rejects_empty_scheme():
with pytest.raises(ValueError):
validate_url("//no-scheme.example.com")
# ---------------------------------------------------------------------------
# safe_fetch - scheme and redirect guards (mocked network)
# ---------------------------------------------------------------------------
def _make_mock_response(content: bytes, status: int = 200):
mock = MagicMock()
mock.__enter__ = lambda s: s
mock.__exit__ = MagicMock(return_value=False)
mock.status = status
mock.code = status
chunks = [content[i:i+65536] for i in range(0, len(content), 65536)] + [b""]
mock.read.side_effect = chunks
return mock
def test_safe_fetch_rejects_file_url():
with pytest.raises(ValueError, match="file"):
safe_fetch("file:///etc/passwd")
def test_safe_fetch_rejects_ftp_url():
with pytest.raises(ValueError, match="ftp"):
safe_fetch("ftp://example.com/file.zip")
def test_safe_fetch_returns_bytes(tmp_path):
mock_resp = _make_mock_response(b"hello world")
with patch("graphify.security._build_opener") as mock_opener_fn:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_opener_fn.return_value = mock_opener
result = safe_fetch("https://example.com/")
assert result == b"hello world"
def test_safe_fetch_raises_on_non_2xx():
mock_resp = _make_mock_response(b"Not Found", status=404)
with patch("graphify.security._build_opener") as mock_opener_fn:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_opener_fn.return_value = mock_opener
with pytest.raises(urllib.error.HTTPError):
safe_fetch("https://example.com/missing")
def test_safe_fetch_raises_on_size_exceeded():
# Build a response larger than max_bytes
big_chunk = b"x" * 65_537
mock_resp = MagicMock()
mock_resp.__enter__ = lambda s: s
mock_resp.__exit__ = MagicMock(return_value=False)
mock_resp.status = 200
mock_resp.code = 200
# Return the chunk twice so total > max_bytes=65536
mock_resp.read.side_effect = [big_chunk, big_chunk, b""]
with patch("graphify.security._build_opener") as mock_opener_fn:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_opener_fn.return_value = mock_opener
with pytest.raises(OSError, match="size limit"):
safe_fetch("https://example.com/huge", max_bytes=65_536)
# ---------------------------------------------------------------------------
# safe_fetch_text
# ---------------------------------------------------------------------------
def test_safe_fetch_text_decodes_utf8():
content = "héllo wörld".encode("utf-8")
mock_resp = _make_mock_response(content)
with patch("graphify.security._build_opener") as mock_opener_fn:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_opener_fn.return_value = mock_opener
result = safe_fetch_text("https://example.com/")
assert result == "héllo wörld"
def test_safe_fetch_text_replaces_bad_bytes():
bad = b"hello \xff world"
mock_resp = _make_mock_response(bad)
with patch("graphify.security._build_opener") as mock_opener_fn:
mock_opener = MagicMock()
mock_opener.open.return_value = mock_resp
mock_opener_fn.return_value = mock_opener
result = safe_fetch_text("https://example.com/")
assert "hello" in result
assert "world" in result
assert "\xff" not in result
# ---------------------------------------------------------------------------
# validate_graph_path
# ---------------------------------------------------------------------------
def test_validate_graph_path_allows_inside_base(tmp_path):
base = tmp_path / "graphify-out"
base.mkdir()
graph = base / "graph.json"
graph.write_text("{}")
result = validate_graph_path(str(graph), base=base)
assert result == graph.resolve()
def test_validate_graph_path_blocks_traversal(tmp_path):
base = tmp_path / "graphify-out"
base.mkdir()
evil = tmp_path / "graphify-out" / ".." / "etc_passwd"
with pytest.raises(ValueError, match="escapes"):
validate_graph_path(str(evil), base=base)
def test_validate_graph_path_requires_base_exists(tmp_path):
base = tmp_path / "graphify-out" # not created
with pytest.raises(ValueError, match="does not exist"):
validate_graph_path(str(base / "graph.json"), base=base)
def test_validate_graph_path_raises_if_file_missing(tmp_path):
base = tmp_path / "graphify-out"
base.mkdir()
with pytest.raises(FileNotFoundError):
validate_graph_path(str(base / "missing.json"), base=base)
# ---------------------------------------------------------------------------
# sanitize_label
# ---------------------------------------------------------------------------
def test_sanitize_label_escapes_html():
assert "<script>" in sanitize_label("<script>")
assert "&" in sanitize_label("foo & bar")
def test_sanitize_label_strips_control_chars():
result = sanitize_label("hello\x00\x1fworld")
assert "\x00" not in result
assert "\x1f" not in result
assert "helloworld" in result
def test_sanitize_label_caps_at_256():
long_label = "a" * 300
assert len(sanitize_label(long_label)) <= 256
def test_sanitize_label_safe_passthrough():
assert sanitize_label("MyClass") == "MyClass"
assert sanitize_label("extract_python") == "extract_python"