package applications

import (
	"crypto/rand"
	"crypto/subtle"
	"cwtch.im/tapir"
	"encoding/json"
	"git.openprivacy.ca/openprivacy/libricochet-go/identity"
	"git.openprivacy.ca/openprivacy/libricochet-go/log"
	"git.openprivacy.ca/openprivacy/libricochet-go/utils"
	"golang.org/x/crypto/ed25519"
	"golang.org/x/crypto/sha3"
	"time"
)

// AuthMessage is exchanged between peers to obtain the Auth Capability
type AuthMessage struct {
	LongTermPublicKey  ed25519.PublicKey
	EphemeralPublicKey ed25519.PublicKey
}

// AuthCapability defines the Authentication Capability granted by AuthApp
const AuthCapability = "AUTH"

// AuthApp is the concrete Application type that handles Authentication
type AuthApp struct {
}

// NewInstance creates a new instance of the AuthApp
func (ea AuthApp) NewInstance() tapir.Application {
	return new(AuthApp)
}

// Init runs the entire AuthApp protocol, at the end of the protocol either the connection is granted AUTH capability
// or the connection is closed.
func (ea AuthApp) Init(connection *tapir.Connection) {
	longTermPubKey := ed25519.PublicKey(connection.ID.PublicKeyBytes())
	epk, esk, _ := ed25519.GenerateKey(rand.Reader)
	ephemeralPublicKey := ed25519.PublicKey(epk)
	ephemeralPrivateKey := ed25519.PrivateKey(esk)
	ephemeralIdentity := identity.InitializeV3("", &ephemeralPrivateKey, &ephemeralPublicKey)
	authMessage := AuthMessage{LongTermPublicKey: longTermPubKey, EphemeralPublicKey: ephemeralPublicKey}
	serialized, _ := json.Marshal(authMessage)
	connection.Send(serialized)
	message := connection.Expect()

	var remoteAuthMessage AuthMessage
	err := json.Unmarshal(message, &remoteAuthMessage)
	if err != nil {
		connection.Close()
		return
	}

	// 3DH Handshake
	l2e := connection.ID.EDH(remoteAuthMessage.EphemeralPublicKey)
	e2l := ephemeralIdentity.EDH(remoteAuthMessage.LongTermPublicKey)
	e2e := ephemeralIdentity.EDH(remoteAuthMessage.EphemeralPublicKey)

	// We need to define an order for the result concatenation so that both sides derive the same key.
	var result [96]byte
	if connection.Outbound {
		copy(result[0:32], l2e)
		copy(result[32:64], e2l)
		copy(result[64:96], e2e)
	} else {
		copy(result[0:32], e2l)
		copy(result[32:64], l2e)
		copy(result[64:96], e2e)
	}
	connection.SetEncryptionKey(sha3.Sum256(result[:]))

	// Wait to Sync
	time.Sleep(time.Second)

	// TODO: Replace this with proper transcript
	challengeRemote, err := json.Marshal(remoteAuthMessage)
	challengeLocal, err := json.Marshal(authMessage)
	challenge := sha3.New512()

	if connection.Outbound {
		challenge.Write(challengeLocal)
		challenge.Write(challengeRemote)
	} else {
		challenge.Write(challengeRemote)
		challenge.Write(challengeLocal)
	}

	// Since we have set the encryption key on the connection the connection will encrypt any messages we send with that key
	// To test that the remote peer has done the same we calculate a challenge hash based on the transcript so far and send it to them
	// We expect the remote to do the same, and compare the two.
	// If successful we extend our auth capability to the connection and reassert the hostname.
	challengeBytes := challenge.Sum([]byte{})
	connection.Send(challengeBytes)
	remoteChallenge := connection.Expect()
	if subtle.ConstantTimeCompare(challengeBytes, remoteChallenge) == 1 {
		connection.SetHostname(utils.GetTorV3Hostname(remoteAuthMessage.LongTermPublicKey))
		connection.SetCapability(AuthCapability)
	} else {
		log.Errorf("Failed Decrypt Challenge: [%x] [%x]\n", remoteChallenge, challengeBytes)
		connection.Close()
	}
}