clock-curve-math 1.1.3

#!/usr/bin/env python3
"""
SageMath implementations of discrete logarithm algorithms.

This module provides correct implementations of discrete logarithm
algorithms that can be used to verify and guide the Rust implementations.
"""

import sage.all as sage
from sage.groups.generic import order_from_multiple
import math


class DiscreteLogAlgorithms:
    """Discrete logarithm algorithm implementations using SageMath."""

    def __init__(self, p, g=None):
        """
        Initialize discrete log setup.

        Args:
            p: Prime modulus
            g: Generator (optional, will find one if not provided)
        """
        self.p = sage.Integer(p)
        self.Fp = sage.Integers(self.p)

        if g is None:
            # Find a generator
            self.g = self.find_generator()
        else:
            self.g = self.Fp(g)

        print(f"Using generator g = {self.g}")

    def find_generator(self):
        """Find a generator of the multiplicative group Fp*."""
        for g in range(2, self.p):
            if self.is_generator(g):
                return self.Fp(g)
        raise ValueError("No generator found")

    def is_generator(self, g):
        """Check if g generates the multiplicative group."""
        g = self.Fp(g)
        if g == 1:
            return False

        # Check that g^{(p-1)/q} != 1 for all prime factors q of p-1
        p_minus_1 = self.p - 1
        factors = p_minus_1.factor()

        for q, _ in factors:
            exponent = p_minus_1 // q
            if g**exponent == 1:
                return False

        return True

    def baby_step_giant_step(self, g, h, max_order=None):
        """
        Baby-step giant-step discrete logarithm algorithm.

        Solves g^x ≡ h mod p

        Args:
            g: Generator
            h: Target value
            max_order: Maximum group order to consider

        Returns:
            x such that g^x ≡ h, or None if not found
        """
        g = self.Fp(g)
        h = self.Fp(h)

        if h == 1:
            return self.Fp(0)  # g^0 = 1

        if max_order is None:
            # Use a reasonable default based on p
            max_order = min(2**20, self.p - 1)  # Up to 1M operations

        m = sage.Integer(math.ceil(math.sqrt(max_order)))

        # Baby steps: compute g^0, g^1, ..., g^{m-1}
        baby_steps = {}
        current = self.Fp(1)
        for i in range(m):
            baby_steps[current] = i
            current = current * g

        # Compute g^{-m}
        g_m = g**m
        g_inv_m = g_m**(-1)

        # Giant steps: compute h * (g^{-m})^j for j = 0,1,2,...
        giant_current = h
        for j in range(m):
            if giant_current in baby_steps:
                return baby_steps[giant_current] + j * m
            giant_current = giant_current * g_inv_m

        return None  # Not found within the search space

    def pollard_rho_discrete_log(self, g, h, max_iterations=10000):
        """
        Pollard's rho algorithm for discrete logarithm.

        This is a simplified implementation that may not work for all cases.
        For now, fall back to baby-step giant-step for reliability.

        Args:
            g: Generator
            h: Target value
            max_iterations: Maximum iterations before giving up

        Returns:
            x such that g^x ≡ h, or None if not found
        """
        # For now, use baby-step giant-step as it's more reliable
        return self.baby_step_giant_step(g, h)

    def index_calculus(self, g, h, factor_base_size=50):
        """
        Index calculus algorithm for discrete logarithm.

        Most efficient for large prime fields.

        Args:
            g: Generator
            h: Target value
            factor_base_size: Size of the factor base

        Returns:
            x such that g^x ≡ h, or None if not found
        """
        # This is a simplified implementation - full index calculus is complex
        # For now, fall back to baby-step giant-step
        return self.baby_step_giant_step(g, h)

    def compute_discrete_log(self, g, h, algorithm='baby_step_giant_step'):
        """
        Compute discrete logarithm using the specified algorithm.

        Args:
            g: Generator
            h: Target value
            algorithm: Algorithm to use ('baby_step_giant_step', 'pollard_rho', 'index_calculus')

        Returns:
            x such that g^x ≡ h
        """
        g = self.Fp(g)
        h = self.Fp(h)

        if algorithm == 'baby_step_giant_step':
            return self.baby_step_giant_step(g, h)
        elif algorithm == 'pollard_rho':
            return self.pollard_rho_discrete_log(g, h)
        elif algorithm == 'index_calculus':
            return self.index_calculus(g, h)
        else:
            raise ValueError(f"Unknown algorithm: {algorithm}")

    def verify_discrete_log(self, g, h, x):
        """Verify that g^x ≡ h."""
        g = self.Fp(g)
        h = self.Fp(h)
        x = sage.Integer(x)

        return g**x == h

    def test_discrete_log(self):
        """Test discrete logarithm implementations."""
        print("Testing discrete logarithm algorithms...")

        # Test with small exponent first
        x_test = 42
        h = self.g ** x_test

        print(f"Testing g^{x_test} = {h}")

        # Test baby-step giant-step
        try:
            result_bsgs = self.baby_step_giant_step(self.g, h)
            print(f"Baby-step giant-step: g^{result_bsgs} = {self.g**result_bsgs}")
            print(f"Correct: {result_bsgs == x_test}")
        except Exception as e:
            print(f"Baby-step giant-step failed: {e}")

        # Test Pollard rho
        try:
            result_rho = self.pollard_rho_discrete_log(self.g, h)
            if result_rho is not None:
                print(f"Pollard rho: g^{result_rho} = {self.g**result_rho}")
                print(f"Correct: {result_rho == x_test}")
            else:
                print("Pollard rho: Not found")
        except Exception as e:
            print(f"Pollard rho failed: {e}")

    def benchmark_algorithms(self, test_cases=None):
        """Benchmark different discrete log algorithms."""
        if test_cases is None:
            test_cases = [10, 100, 1000, 10000]

        print("Benchmarking discrete logarithm algorithms...")
        print("Exponent | BSG S (ms) | Pollard Rho (ms)")
        print("-" * 40)

        import time

        for x in test_cases:
            h = self.g ** x

            # Benchmark BSG S
            start = time.time()
            result_bsgs = self.baby_step_giant_step(self.g, h)
            time_bsgs = (time.time() - start) * 1000

            # Benchmark Pollard rho
            start = time.time()
            result_rho = self.pollard_rho_discrete_log(self.g, h)
            time_rho = (time.time() - start) * 1000

            print(">8")

            if result_bsgs != x or (result_rho is not None and result_rho != x):
                print("  WARNING: Incorrect result!")


if __name__ == "__main__":
    # Test with a small prime for demonstration
    p = 101  # Small prime for testing
    dl = DiscreteLogAlgorithms(p)
    dl.test_discrete_log()

    # Benchmark with larger prime but small exponents
    p_large = 2**31 - 1  # Mersenne prime
    dl_large = DiscreteLogAlgorithms(p_large)
    dl_large.benchmark_algorithms([10, 100, 1000])