mathtools/cmd/discreteLog.go

/*
Copyright © 2025 filifa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package cmd

import (
	"errors"
	"fmt"
	"math/big"

	"github.com/spf13/cobra"
)

var discreteLogModulus string
var discreteLogBase string
var discreteLogElement string
var discreteLogOrder string

// whyyyy doesn't math/big have a ceil functionnnn
func ceilSqrt(x *big.Int) *big.Int {
	z := new(big.Int).Sqrt(x)
	s := new(big.Int).Exp(z, big.NewInt(2), nil)
	if s.Cmp(x) != 0 {
		z.Add(z, big.NewInt(1))
	}

	return z
}

// TODO: this can be extended to work with n, b not coprime
// https://cp-algorithms.com/algebra/discrete-log.html
func babyStepGiantStep(n, b, x, order *big.Int) (*big.Int, error) {
	z := new(big.Int).GCD(nil, nil, b, n)
	if z.Cmp(big.NewInt(1)) != 0 {
		return nil, fmt.Errorf("base %v and modulus %v are not coprime", b, n)
	}

	var m *big.Int
	if order == nil {
		// m = ceil(sqrt(n - 1))
		z := big.NewInt(1)
		z.Sub(n, z)
		m = ceilSqrt(z)
	} else {
		m = ceilSqrt(order)
	}

	table := make(map[string]*big.Int)
	for j := big.NewInt(1); j.Cmp(m) <= 0; j.Add(j, big.NewInt(1)) {
		a := new(big.Int).Exp(b, j, n)
		table[a.String()] = new(big.Int).Set(j)
	}

	// p = b^-m modulo n
	p := new(big.Int).Neg(m)
	p.Exp(b, p, n)

	gamma := new(big.Int).Set(x)

	for i := big.NewInt(0); i.Cmp(m) == -1; i.Add(i, big.NewInt(1)) {
		j, ok := table[gamma.String()]
		if ok {
			i.Mul(i, m)
			i.Add(i, j)
			return i, nil
		}

		gamma.Mul(gamma, p)
		gamma.Mod(gamma, n)
	}

	return nil, errors.New("no solution")
}

func discreteLog(cmd *cobra.Command, args []string) {
	m, ok := new(big.Int).SetString(discreteLogModulus, 10)
	if !ok {
		cobra.CheckErr("invalid modulus " + discreteLogModulus)
	}

	b, ok := new(big.Int).SetString(discreteLogBase, 10)
	if !ok {
		cobra.CheckErr("invalid base " + discreteLogBase)
	}

	x, ok := new(big.Int).SetString(discreteLogElement, 10)
	if !ok {
		cobra.CheckErr("invalid element " + discreteLogElement)
	}

	var order *big.Int
	if discreteLogOrder != "" {
		order, ok = new(big.Int).SetString(discreteLogOrder, 10)
		if !ok {
			cobra.CheckErr("invalid order " + discreteLogOrder)
		}
	}

	k, err := babyStepGiantStep(m, b, x, order)
	if err != nil {
		cobra.CheckErr(err)
	}

	fmt.Println(k)
}

// discreteLogCmd represents the discreteLog command
var discreteLogCmd = &cobra.Command{
	Use:   "discrete-log -b N -m N -e N",
	Short: "Compute the discrete logarithm",
	Long: `Compute the discrete logarithm.

Given a base b, modulus m, and element e, compute a value k such that b^k = e (mod m).

Note that no efficient method of finding the discrete logarithm is currently known. For slightly improved performance, the order of the group (i.e. the totient of m) can be provided.`,
	Run: discreteLog,
}

func init() {
	rootCmd.AddCommand(discreteLogCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// discreteLogCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// discreteLogCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

	discreteLogCmd.Flags().StringVarP(&discreteLogModulus, "modulus", "m", "", "modulus of the cyclic group")
	discreteLogCmd.MarkFlagRequired("modulus")

	discreteLogCmd.Flags().StringVarP(&discreteLogBase, "base", "b", "", "logarithmic base")
	discreteLogCmd.MarkFlagRequired("base")

	discreteLogCmd.Flags().StringVarP(&discreteLogElement, "element", "e", "", "element of the cyclic group to compute logarithm of")
	discreteLogCmd.MarkFlagRequired("element")

	discreteLogCmd.Flags().StringVar(&discreteLogOrder, "order", "", "order of the cyclic group")
}