flaskpaste/tests/test_pow.py

"""Tests for proof-of-work system."""

import hashlib
import json
import time


class TestPowChallenge:
    """Tests for GET /challenge endpoint."""

    def test_challenge_disabled_by_default_in_testing(self, client):
        """Challenge endpoint returns disabled when POW_DIFFICULTY=0."""
        response = client.get("/challenge")
        assert response.status_code == 200
        data = json.loads(response.data)
        assert data["enabled"] is False
        assert data["difficulty"] == 0

    def test_challenge_returns_token_when_enabled(self, app, client):
        """Challenge endpoint returns valid token when enabled."""
        app.config["POW_DIFFICULTY"] = 8

        response = client.get("/challenge")
        assert response.status_code == 200
        data = json.loads(response.data)

        assert data["enabled"] is True
        assert data["difficulty"] == 8
        assert "nonce" in data
        assert "expires" in data
        assert "token" in data
        assert len(data["nonce"]) == 32  # 16 bytes hex
        assert data["expires"] > int(time.time())


class TestPowVerification:
    """Tests for PoW verification in paste creation."""

    def test_paste_without_pow_when_disabled(self, client, sample_text):
        """Paste creation works without PoW when disabled."""
        response = client.post("/", data=sample_text, content_type="text/plain")
        assert response.status_code == 201

    def test_paste_requires_pow_when_enabled(self, app, client, sample_text):
        """Paste creation fails without PoW when enabled."""
        app.config["POW_DIFFICULTY"] = 8

        response = client.post("/", data=sample_text, content_type="text/plain")
        assert response.status_code == 400
        data = json.loads(response.data)
        assert "Proof-of-work required" in data["error"]

    def test_paste_with_valid_pow(self, app, client, sample_text):
        """Paste creation succeeds with valid PoW."""
        app.config["POW_DIFFICULTY"] = 8

        # Get challenge
        ch_response = client.get("/challenge")
        ch_data = json.loads(ch_response.data)

        # Solve PoW
        nonce = ch_data["nonce"]
        difficulty = ch_data["difficulty"]
        solution = solve_pow(nonce, difficulty)

        # Submit with solution
        response = client.post(
            "/",
            data=sample_text,
            content_type="text/plain",
            headers={
                "X-PoW-Token": ch_data["token"],
                "X-PoW-Solution": str(solution),
            },
        )
        assert response.status_code == 201

    def test_paste_with_invalid_solution(self, app, client, sample_text):
        """Paste creation fails with wrong solution."""
        app.config["POW_DIFFICULTY"] = 8

        # Get challenge
        ch_response = client.get("/challenge")
        ch_data = json.loads(ch_response.data)

        # Submit with wrong solution
        response = client.post(
            "/",
            data=sample_text,
            content_type="text/plain",
            headers={
                "X-PoW-Token": ch_data["token"],
                "X-PoW-Solution": "999999999",
            },
        )
        assert response.status_code == 400
        data = json.loads(response.data)
        assert "Insufficient work" in data["error"]

    def test_paste_with_expired_challenge(self, app, client, sample_text):
        """Paste creation fails with expired challenge."""
        app.config["POW_DIFFICULTY"] = 8
        app.config["POW_CHALLENGE_TTL"] = 1  # 1 second

        # Get challenge
        ch_response = client.get("/challenge")
        ch_data = json.loads(ch_response.data)

        # Solve PoW
        solution = solve_pow(ch_data["nonce"], ch_data["difficulty"])

        # Wait for expiry
        time.sleep(2)

        # Submit with expired challenge
        response = client.post(
            "/",
            data=sample_text,
            content_type="text/plain",
            headers={
                "X-PoW-Token": ch_data["token"],
                "X-PoW-Solution": str(solution),
            },
        )
        assert response.status_code == 400
        data = json.loads(response.data)
        assert "expired" in data["error"].lower()

    def test_paste_with_tampered_token(self, app, client, sample_text):
        """Paste creation fails with tampered token."""
        app.config["POW_DIFFICULTY"] = 8

        # Get challenge
        ch_response = client.get("/challenge")
        ch_data = json.loads(ch_response.data)

        # Tamper with token (change difficulty)
        parts = ch_data["token"].split(":")
        parts[2] = "1"  # Lower difficulty
        tampered_token = ":".join(parts)

        # Submit with tampered token
        response = client.post(
            "/",
            data=sample_text,
            content_type="text/plain",
            headers={
                "X-PoW-Token": tampered_token,
                "X-PoW-Solution": "0",
            },
        )
        assert response.status_code == 400
        data = json.loads(response.data)
        assert "signature" in data["error"].lower() or "mismatch" in data["error"].lower()


def solve_pow(nonce, difficulty):
    """Solve proof-of-work challenge (test helper)."""
    n = 0
    while True:
        work = f"{nonce}:{n}".encode()
        hash_bytes = hashlib.sha256(work).digest()

        zero_bits = 0
        for byte in hash_bytes:
            if byte == 0:
                zero_bits += 8
            else:
                zero_bits += 8 - byte.bit_length()
                break

        if zero_bits >= difficulty:
            return n
        n += 1