// Package rocacheck checks if a key was generated by broken Infineon code and
// is vulnerable to factorization via the Return of Coppersmith's Attack (ROCA)
// / CVE-2017-15361.
package rocacheck

import (
	"crypto/rsa"
	"math/big"
)

type test struct {
	Prime        *big.Int
	Fingerprints map[int64]struct{}
}

var tests = make([]test, 17)

func init() {
	bigOne := big.NewInt(1)
	n := &big.Int{}
	// relations table from https://github.com/crocs-muni/roca/pull/40
	for i, r := range [][2]int64{
		{2, 11}, {6, 13}, {8, 17}, {9, 19}, {3, 37}, {26, 53}, {20, 61},
		{35, 71}, {24, 73}, {13, 79}, {6, 97}, {51, 103}, {53, 107},
		{54, 109}, {42, 127}, {50, 151}, {78, 157},
	} {
		fps := make(map[int64]struct{})
		bp := big.NewInt(r[1])
		br := big.NewInt(r[0])
		for j := int64(0); j < r[1]; j++ {
			if n.Exp(big.NewInt(j), br, bp).Cmp(bigOne) == 0 {
				fps[j] = struct{}{}
			}
		}
		tests[i] = test{
			Prime:        big.NewInt(r[1]),
			Fingerprints: fps,
		}
	}
}

// IsWeak returns true if a RSA public key is vulnerable to Return of
// Coppersmith's Attack (ROCA).
func IsWeak(k *rsa.PublicKey) bool {
	tmp := &big.Int{}
	for _, t := range tests {
		if _, ok := t.Fingerprints[tmp.Mod(k.N, t.Prime).Int64()]; !ok {
			return false
		}
	}
	return true
}