190 lines
7.9 KiB
Python
190 lines
7.9 KiB
Python
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import pwd
|
|
import socket
|
|
import struct
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
|
|
from PIL import Image
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
|
|
from auth import AuthFlow
|
|
from greetd_ipc import Client, ProtocolError, SESSION_COMMAND, MAX_FRAME
|
|
from accounts import users
|
|
from system_state import battery, backlight, set_brightness
|
|
from avatar_backend import AvatarError, render_png, source_path
|
|
from avatar_icon import system_icon
|
|
|
|
SECRET = "LABFY_SYNTHETIC_SECRET_59f3c4e2"
|
|
|
|
|
|
class FakeGreetd:
|
|
def __init__(self, replies):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.path = self.temp.name + "/socket"
|
|
self.server = socket.socket(socket.AF_UNIX)
|
|
self.server.bind(self.path)
|
|
self.server.listen(1)
|
|
self.replies = replies
|
|
self.requests = []
|
|
self.thread = threading.Thread(target=self.run, daemon=True)
|
|
self.thread.start()
|
|
|
|
def run(self):
|
|
connection, _ = self.server.accept()
|
|
with connection:
|
|
for reply in self.replies:
|
|
header = connection.recv(4)
|
|
if not header:
|
|
break
|
|
size = struct.unpack("=I", header)[0]
|
|
assert size <= MAX_FRAME
|
|
data = b""
|
|
while len(data) < size:
|
|
data += connection.recv(size - len(data))
|
|
request = json.loads(data)
|
|
self.requests.append(request)
|
|
if request["type"] == "start_session":
|
|
assert request["cmd"] == SESSION_COMMAND
|
|
assert request["env"] == []
|
|
encoded = json.dumps(reply).encode()
|
|
connection.sendall(struct.pack("=I", len(encoded)) + encoded)
|
|
|
|
def close(self):
|
|
self.server.close()
|
|
self.thread.join(2)
|
|
self.temp.cleanup()
|
|
|
|
|
|
def msg(kind, value="Message PAM"):
|
|
return {"type": "auth_message", "auth_message_type": kind, "auth_message": value}
|
|
|
|
|
|
class Tests(unittest.TestCase):
|
|
def test_full_pam_sequence_and_secret_absent_from_outputs(self):
|
|
fake = FakeGreetd([msg("secret"), msg("info"), msg("visible"), msg("error"), msg("secret"),
|
|
{"type": "success"}, {"type": "success"}])
|
|
output = io.StringIO()
|
|
try:
|
|
with contextlib.redirect_stdout(output), contextlib.redirect_stderr(output):
|
|
flow = AuthFlow(fake.path)
|
|
self.assertEqual(flow.begin("example")["auth_message_type"], "secret")
|
|
self.assertEqual(flow.answer(SECRET)["auth_message_type"], "info")
|
|
self.assertEqual(flow.answer()["auth_message_type"], "visible")
|
|
self.assertEqual(flow.answer("code")["auth_message_type"], "error")
|
|
self.assertEqual(flow.answer()["auth_message_type"], "secret")
|
|
self.assertEqual(flow.answer("second")["type"], "success")
|
|
self.assertEqual(flow.start()["type"], "success")
|
|
self.assertNotIn(SECRET, output.getvalue())
|
|
self.assertEqual(fake.requests[1]["response"], SECRET)
|
|
self.assertNotIn("response", fake.requests[2])
|
|
self.assertEqual(fake.requests[-1]["cmd"], SESSION_COMMAND)
|
|
finally:
|
|
fake.close()
|
|
|
|
def test_auth_error_retry(self):
|
|
first = FakeGreetd([msg("secret"), {"type": "error", "error_type": "auth_error", "description": "bad"}, {"type": "success"}])
|
|
second = FakeGreetd([msg("visible"), {"type": "success"}, {"type": "success"}])
|
|
try:
|
|
flow = AuthFlow(first.path)
|
|
flow.begin("example")
|
|
self.assertEqual(flow.answer("wrong")["error_type"], "auth_error")
|
|
self.assertFalse(flow.authenticated)
|
|
flow.socket_path = second.path
|
|
self.assertEqual(flow.begin("manual-nss-user")["auth_message_type"], "visible")
|
|
self.assertEqual(flow.answer("code")["type"], "success")
|
|
self.assertEqual(flow.start()["type"], "success")
|
|
self.assertEqual(second.requests[0]["username"], "manual-nss-user")
|
|
finally:
|
|
first.close()
|
|
second.close()
|
|
|
|
def test_accounts(self):
|
|
Row = pwd.struct_passwd
|
|
entries = [Row(("system", "x", 999, 1, "", "/", "/bin/bash")),
|
|
Row(("human", "x", 1000, 1, "", "/", "/bin/bash")),
|
|
Row(("blocked", "x", 1001, 1, "", "/", "/usr/bin/nologin")),
|
|
Row(("false", "x", 1002, 1, "", "/", "/bin/false")),
|
|
Row(("another", "x", 1003, 1, "Example", "/", "/bin/zsh"))]
|
|
self.assertEqual(users(entries, (1000, 60000)), ["another", "human"])
|
|
|
|
def test_battery_and_backlight(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
root = Path(temp)
|
|
self.assertIsNone(battery(root))
|
|
bat = root / "BAT1"
|
|
bat.mkdir()
|
|
(bat / "capacity").write_text("60")
|
|
for status in ("Charging", "Discharging", "Full", "Not charging"):
|
|
(bat / "status").write_text(status)
|
|
self.assertEqual(battery(root), (60, status))
|
|
light = root / "backlight" / "amdgpu_bl1"
|
|
light.mkdir(parents=True)
|
|
(light / "max_brightness").write_text("400000")
|
|
(light / "brightness").write_text("200000")
|
|
self.assertEqual(backlight(root / "backlight"), (light, 400000, 200000))
|
|
set_brightness(light, 400000, 10)
|
|
self.assertEqual((light / "brightness").read_text(), "40000")
|
|
set_brightness(light, 400000, 100)
|
|
self.assertEqual((light / "brightness").read_text(), "400000")
|
|
(light / "max_brightness").write_text("0")
|
|
self.assertIsNone(backlight(root / "backlight"))
|
|
(light / "max_brightness").unlink()
|
|
self.assertIsNone(backlight(root / "backlight"))
|
|
self.assertIsNone(backlight(root / "missing"))
|
|
with self.assertRaises(OSError):
|
|
set_brightness(root / "missing", 400000, 50)
|
|
|
|
def test_reject_oversized_response(self):
|
|
temp = tempfile.TemporaryDirectory()
|
|
path = temp.name + "/socket"
|
|
server = socket.socket(socket.AF_UNIX)
|
|
server.bind(path)
|
|
server.listen(1)
|
|
def serve():
|
|
conn, _ = server.accept()
|
|
with conn:
|
|
conn.recv(4096)
|
|
conn.sendall(struct.pack("=I", MAX_FRAME + 1))
|
|
thread = threading.Thread(target=serve, daemon=True)
|
|
thread.start()
|
|
try:
|
|
with Client(path) as client:
|
|
with self.assertRaises(ProtocolError):
|
|
client.create("example")
|
|
finally:
|
|
server.close()
|
|
thread.join(2)
|
|
temp.cleanup()
|
|
|
|
def test_avatar_conversion_and_system_boundary(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
root = Path(temp)
|
|
original = root / "portrait.jpg"
|
|
Image.new("RGB", (400, 200), (10, 20, 30)).save(original)
|
|
self.assertEqual(source_path(original.as_uri()), original)
|
|
converted = root / "avatar.png"
|
|
render_png(original, converted)
|
|
with Image.open(converted) as image:
|
|
self.assertEqual(image.format, "PNG")
|
|
self.assertEqual(image.size, (256, 256))
|
|
with self.assertRaises(AvatarError):
|
|
source_path("https://example.invalid/avatar.png")
|
|
too_large = root / "too-large.jpg"
|
|
too_large.write_bytes(b"x" * (12 * 1024 * 1024 + 1))
|
|
with self.assertRaises(AvatarError):
|
|
source_path(str(too_large))
|
|
self.assertIsNone(system_icon("../etc/passwd", root))
|
|
self.assertIsNone(system_icon("example", root))
|
|
(root / "example").symlink_to(converted)
|
|
self.assertIsNone(system_icon("example", root))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|