package spam

import (
	"crypto/rand"
	"crypto/sha256"
	"cwtch.im/cwtch/protocol"
	"git.openprivacy.ca/openprivacy/libricochet-go/utils"
	"git.openprivacy.ca/openprivacy/libricochet-go/wire/control"
	"github.com/golang/protobuf/proto"
	"io"
	//"fmt"
)

// Guard implements a spam protection mechanism for Cwtch Servers.
type Guard struct {
	Difficulty int
	nonce      [24]byte
}

func getRandomness(arr *[24]byte) {
	if _, err := io.ReadFull(rand.Reader, arr[:]); err != nil {
		utils.CheckError(err)
	}
}

//GenerateChallenge returns a channel result packet with a spamguard challenge nonce
func (sg *Guard) GenerateChallenge(channelID int32) []byte {

	cr := &Protocol_Data_Control.ChannelResult{
		ChannelIdentifier: proto.Int32(channelID),
		Opened:            proto.Bool(true),
	}

	var nonce [24]byte
	getRandomness(&nonce)
	sg.nonce = nonce
	err := proto.SetExtension(cr, protocol.E_ServerNonce, sg.nonce[:])
	utils.CheckError(err)

	pc := &Protocol_Data_Control.Packet{
		ChannelResult: cr,
	}
	ret, err := proto.Marshal(pc)
	utils.CheckError(err)
	return ret
}

// SolveChallenge takes in a challenge and a message and returns a solution
// The solution is a 24 byte nonce which when hashed with the challenge and the message
// produces a sha256 hash with Difficulty leading 0s
func (sg *Guard) SolveChallenge(challenge []byte, message []byte) []byte {
	solved := false
	var spamguard [24]byte
	sum := sha256.Sum256([]byte{})
	solve := make([]byte, len(challenge)+len(message)+len(spamguard))
	for !solved {

		getRandomness(&spamguard)

		copy(solve[0:], challenge[:])
		copy(solve[len(challenge):], message[:])
		copy(solve[len(challenge)+len(message):], spamguard[:])

		sum = sha256.Sum256(solve)

		solved = true
		for i := 0; i < sg.Difficulty; i++ {
			if sum[i] != 0x00 {
				solved = false
			}
		}
	}
	//fmt.Printf("[SOLVED] %x\n",sha256.Sum256(solve))
	return spamguard[:]
}

// ValidateChallenge returns true if the message and spamguard pass the challenge
func (sg *Guard) ValidateChallenge(message []byte, spamguard []byte) bool {
	if len(spamguard) != 24 {
		return false
	}

	// If the message is too large just throw it away.
	if len(message) > 2048 {
		return false
	}

	solve := make([]byte, len(sg.nonce)+len(message)+len(spamguard))
	copy(solve[0:], sg.nonce[:])
	copy(solve[len(sg.nonce):], message[:])
	copy(solve[len(sg.nonce)+len(message):], spamguard[:])
	sum := sha256.Sum256(solve)

	for i := 0; i < sg.Difficulty; i++ {
		if sum[i] != 0x00 {
			return false
		}
	}
	return true
}