package main

import (
	"bytes"
	"errors"
)

type Winbox struct {
	task  *Task
	conn  *Connection
	stage int

	user, pass                    string
	w                             *WCurve
	sa, xwa, xwb, j, z, secret, i []byte
	clientCC, serverCC, msg, resp []byte
	xwaParity, xwbParity          bool
}

func NewWinbox(task *Task, conn *Connection) *Winbox {
	winbox := Winbox{}
	winbox.task = task
	winbox.conn = conn
	winbox.xwaParity = false
	winbox.xwbParity = false
	winbox.w = newWCurve()
	winbox.stage = -1
	winbox.user = task.login
	winbox.pass = task.password
	return &winbox
}

func (winbox *Winbox) genSharedSecret(salt []byte) error {
	winbox.i = genPasswordValidatorPriv(winbox.user, winbox.pass, salt)
	xGamma, _ := winbox.w.genPublicKey(winbox.i)
	v := winbox.w.redp1(xGamma, true)

	wb := winbox.w.liftX(NewBigintFromBytes(winbox.xwb), winbox.xwbParity)
	if wb == nil {
		winbox.stage = -1
		return errors.New("liftX failed")
	}

	wb = wb.Add(v)

	xwaCombined := append(winbox.xwa, winbox.xwb...)
	winbox.j = getSHA2Digest(xwaCombined)
	pt := NewBigintFromBytes(winbox.i).Mul(NewBigintFromBytes(winbox.j))
	pt = pt.Add(NewBigintFromBytes(winbox.sa))
	pt = winbox.w.finiteFieldValue(pt)

	mp := wb.Mul(pt)

	winbox.z, _ = winbox.w.toMontgomery(mp)
	winbox.secret = getSHA2Digest(winbox.z)

	return nil
}

func (winbox *Winbox) publicKeyExchange() {
	winbox.sa, _ = genRandomBytes(32)
	winbox.xwa, winbox.xwaParity = winbox.w.genPublicKey(winbox.sa)

	lx := winbox.w.liftX(NewBigintFromBytes(winbox.xwa), winbox.xwaParity)
	if lx == nil {
		log("winbox", 5, "liftX failed in PKE")
		winbox.stage = -1
		return
	}

	if !winbox.w.check(lx) {
		log("winbox", 5, "curve check failed")
		winbox.stage = -1
		return
	}

	winbox.msg = append([]byte(winbox.user), byte(0))
	winbox.msg = append(winbox.msg, winbox.xwa...)
	if winbox.xwaParity {
		winbox.msg = append(winbox.msg, byte(1))
	} else {
		winbox.msg = append(winbox.msg, byte(0))
	}

	header := append([]byte{byte(len(winbox.msg))}, byte(6))
	winbox.msg = append(header, winbox.msg...)
	winbox.stage = 1
}

func (winbox *Winbox) confirmation() error {
	if len(winbox.resp) <= 2 {
		log("winbox", 5, "response size must be greater than 2 (got %v)", len(winbox.resp))
		winbox.stage = -1
		return errors.New("invalid response size")
	}

	respLen := winbox.resp[0]
	winbox.resp = winbox.resp[2:]
	if len(winbox.resp) != int(respLen) {
		log("winbox", 5, "invalid challenge response size: got %v, expected %v", len(winbox.resp), respLen)
		winbox.stage = -1
		return errors.New("invalid challenge response size")
	}

	winbox.xwb = winbox.resp[:32]
	if winbox.resp[32] == 0 {
		winbox.xwbParity = false
	} else {
		winbox.xwbParity = true
	}

	salt := winbox.resp[33:]
	if len(salt) != 16 {
		// this means that there is no such login,
		// or that its an old routeros version
		log("winbox", 5, "invalid salt size: got %v, expected 16", len(salt))

		// report this finding to endpoint manager
		winbox.task.EventWithParm(TH_NoSuchLogin, winbox.user)

		winbox.stage = -1
		return nil // this is not a fatal error for an endpoint
	}

	err := winbox.genSharedSecret(salt)
	if err != nil {
		return err
	}

	winbox.j = getSHA2Digest(append(winbox.xwa, winbox.xwb...))
	winbox.clientCC = getSHA2Digest(append(winbox.j, winbox.z...))

	header := append([]byte{byte(len(winbox.clientCC))}, byte(6))
	winbox.msg = append(header, winbox.clientCC...)
	winbox.stage = 2
	return nil
}

func (winbox *Winbox) sendAndRecv() error {
	if len(winbox.msg) > 0 && winbox.conn != nil {
		err := winbox.conn.Send(winbox.msg)
		winbox.msg = []byte{}

		if err != nil {
			log("winbox", 5, "failed to send: %v", err.Error())
			return err
		}

		winbox.resp, err = winbox.conn.Recv()
		if err != nil {
			log("winbox", 5, "failed to recv: %v", err.Error())
			return err
		}
	}

	return nil
}

func (winbox *Winbox) TryLogin() (result bool, err error) {
	log("winbox", 4, "login: stage 1, PKE")
	winbox.publicKeyExchange()

	log("winbox", 4, "login: stage 1, PKE OK, sending")
	err = winbox.sendAndRecv()
	if err != nil {
		return false, err
	}

	log("winbox", 4, "login: stage 2, confirmation")
	err = winbox.confirmation()
	if err != nil {
		return false, err
	}

	if winbox.stage == -1 { // confirmation failed but no error?
		return false, nil // report that its a bad login
	}

	log("winbox", 4, "login: stage 2, confirmation OK, sending")
	err = winbox.sendAndRecv()
	if err != nil {
		return false, err
	}

	log("winbox", 4, "login: stage 3")
	a1 := append(winbox.j, winbox.clientCC...)
	winbox.serverCC = getSHA2Digest(append(a1, winbox.z...))

	return bytes.Equal(winbox.resp[2:], winbox.serverCC), nil
}