package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha512"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/julienschmidt/httprouter"

	"git.nemunai.re/srs/adlin/libadlin"
)

const IPgwDMZ = "172.23.200.1"

type Challenge struct {
	Accessible []func(*adlin.Student, *http.Request) error
	Check      func(*adlin.Student, *givenToken, int) error
}

/* Restrictions */

func noAccessRestriction(*adlin.Student, *http.Request) error {
	return nil
}

func noAccess(*adlin.Student, *http.Request) error {
	return errors.New("This challenge cannot be accessed this way. ")
}

func accessFrom(ip string) func(_ *adlin.Student, r *http.Request) error {
	return func(_ *adlin.Student, r *http.Request) error {
		if r.Header.Get("X-Forwarded-By") != ip {
			return errors.New("This challenge is not accessible this way.")
		}
		return nil
	}
}

func notAccessFrom(ip string) func(_ *adlin.Student, r *http.Request) error {
	return func(_ *adlin.Student, r *http.Request) error {
		if r.Header.Get("X-Forwarded-By") == ip {
			return errors.New("This challenge is not accessible this way.")
		}
		return nil
	}
}

func maxProxy(nb int) func(_ *adlin.Student, r *http.Request) error {
	return func(_ *adlin.Student, r *http.Request) error {
		if len(strings.Split(r.Header.Get("X-Forwarded-For"), ",")) > nb {
			return errors.New("This challenge is not accessible this way.")
		}
		return nil
	}
}

func sslOnly(_ *adlin.Student, r *http.Request) error {
	if r.Header.Get("X-Forwarded-Proto") != "https" {
		return errors.New("This challenge should be performed over TLS.")
	}
	return nil
}

/* Challenges */

func challengeOk(s *adlin.Student, t *givenToken, chid int) error {
	pkey := s.GetPKey()
	if expectedToken, err := GenerateToken(pkey, 0, []byte(t.Data[0])); err != nil {
		return err
	} else if !hmac.Equal(expectedToken, t.token) {
		return errors.New("This is not the expected token.")
	} else {
		return nil
	}
}

func challenge42(s *adlin.Student, t *givenToken, chid int) error {
	pkey := s.GetPKey()
	if expectedToken, err := GenerateToken(pkey, chid, []byte("42")); err != nil {
		return err
	} else if !hmac.Equal(expectedToken, t.token) {
		return errors.New("This is not the expected token.")
	} else {
		return nil
	}
}

func challengeDNS(s *adlin.Student, t *givenToken, chid int) error {
	pkey := s.GetPKey()
	if expectedToken, err := GenerateToken(pkey, chid, []byte("8dde678132d6c558fc6adaeb9f1d53bf6ec7b876308cf98c48604caa9138523c1ce58b672c87c7e7d9b7248b81804d3940dbf20bf263eeb683244f7c1143712d")); err != nil {
		return err
	} else if !hmac.Equal(expectedToken, t.token) {
		return errors.New("This is not the expected token.")
	}

	return nil
}

func challengeTime(s *adlin.Student, t *givenToken, chid int) error {
	pkey := s.GetPKey()
	if expectedToken, err := GenerateToken(pkey, chid, []byte(t.Data[0])); err != nil {
		return err
	} else if !hmac.Equal(expectedToken, t.token) {
		return errors.New("This is not the expected token.")
	} else if t, err := strconv.ParseInt(t.Data[0], 10, 64); err != nil {
		return err
	} else {
		var rt time.Time
		if t > 3000000000 {
			rt = time.Unix(t/1000000000, t%1000000000)
		} else {
			rt = time.Unix(t, 0)
		}
		if rt.After(time.Now()) {
			return errors.New("You seem to live in the future...")
		} else if d := rt.Sub(time.Now()); d.Minutes() < -10 {
			return errors.New(fmt.Sprintf("Your time is %.0f minutes from us", d.Minutes()))
		}
		return nil
	}
}

func challengePing(s *adlin.Student, t *givenToken, chid int) error {
	var expected []byte
	switch s.Id % 5 {
	case 1:
		expected = []byte("baaaaaad")
	case 2:
		expected = []byte("baadfood")
	case 3:
		expected = []byte("baddcafe")
	case 4:
		expected = []byte("cafebabe")
	default:
		expected = []byte("deadbeef")
	}

	// Allow 1 iteration or the 10 ones
	var expected10 []byte
	for i := 0; i < 10; i++ {
		expected10 = append(expected10, expected...)
	}

	// Make it case insensitive
	expectedC := bytes.ToUpper(expected)
	expected10C := bytes.ToUpper(expected10)

	pkey := s.GetPKey()
	if expectedToken, err := GenerateToken(pkey, 0, expected); err != nil {
		return err
	} else if hmac.Equal(expectedToken, t.token) {
		return nil
	} else if expectedToken, err := GenerateToken(pkey, 0, expectedC); err != nil {
		return err
	} else if hmac.Equal(expectedToken, t.token) {
		return nil
	} else if expectedToken, err := GenerateToken(pkey, 0, expected10); err != nil {
		return err
	} else if hmac.Equal(expectedToken, t.token) {
		return nil
	} else if expectedToken, err := GenerateToken(pkey, 0, expected10C); err != nil {
		return err
	} else if hmac.Equal(expectedToken, t.token) {
		return nil
	} else {
		return errors.New("This is not the expected token.")
	}
}

func challengeDisk(s *adlin.Student, t *givenToken, chid int) error {
	pkey := fmt.Sprintf("%x", s.GetPKey())

	n1, err := strconv.Atoi(t.Data[0][0:2])
	if err != nil {
		return err
	}
	n2, err := strconv.Atoi(t.Data[0][2:4])
	if err != nil {
		return err
	}

	sum := make([]byte, hex.DecodedLen(len(t.Data[0][4:])))
	if _, err := hex.Decode(sum, []byte(t.Data[0][4:])); err != nil {
		return err
	}

	if n1+n2 > len(pkey) {
		n2 = len(pkey) - n1
	}
	if n1+n2 < n1 {
		return errors.New("This is not the expected token.")
	}

	expectedToken := sha512.Sum512([]byte(pkey[n1 : n1+n2]))

	if !hmac.Equal(expectedToken[:], sum) {
		return errors.New("This is not the expected token.")
	} else {
		return nil
	}
}

func challengeEMail(s *adlin.Student, t *givenToken, chid int) error {
	return errors.New("This is not the expected token.")
}

func challengeAPIEMail(_ httprouter.Params, body []byte) (interface{}, error) {
	var st map[string]string
	if err := json.Unmarshal(body, &st); err != nil {
		return nil, err
	}

	if std, err := adlin.GetStudentByLogin(st["login"]); err != nil {
		return nil, err
	} else {
		if _, err := std.UnlockNewChallenge(9, st["token"]); err != nil {
			log.Println(err)
			return nil, err
		}

		log.Printf("%s just unlock ch#%d\n", std.Login, 9)
		return "Success", nil
	}
}

var challenges []Challenge

func init() {
	challenges = []Challenge{
		/* Challenge 0 : shadow updated */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccessRestriction},
			Check:      challengeOk,
		},

		/* Challenge 1 : 42 */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccessRestriction},
			Check:      challenge42,
		},

		/* Challenge 2 : 42 from DMZ */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{accessFrom(IPgwDMZ)},
			Check:      challenge42,
		},

		/* Challenge 3 : ssl (+ ntp) */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{accessFrom(IPgwDMZ), sslOnly},
			Check:      challengeTime,
		},

		/* Challenge 4 : DNS TXT */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{accessFrom(IPgwDMZ), sslOnly},
			Check:      challengeDNS,
		},

		/* Challenge 5 : time net */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{maxProxy(1)},
			Check:      challengeTime,
		},

		/* Bonus 0 : toctoc (read in source code) */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccessRestriction},
			Check:      challenge42,
		},

		/* Bonus 1 : echo request */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccessRestriction},
			Check:      challengePing,
		},

		/* Bonus 2 : disk */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccessRestriction},
			Check:      challengeDisk,
		},

		/* Bonus 3 : mail */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccessRestriction},
			Check:      challengeEMail,
		},

		/* wg step */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccessRestriction},
			Check:      challengeOk,
		},

		/* Last : SSH key, see ssh.go:156 in NewKey function */
		Challenge{
			Accessible: []func(*adlin.Student, *http.Request) error{noAccess},
		},
	}

	router.GET("/challenges", apiHandler(getChallengeList))
	router.GET("/challenge/:chid", rawHandler(responseHandler(accessibleChallenge)))
	router.POST("/challenge", rawHandler(responseHandler(challengeHandler(receiveToken))))
	router.POST("/challenge/:chid", rawHandler(responseHandler(receiveChallenge)))
	router.POST("/toctoc", rawHandler(responseHandler(definedChallengeHandler(receiveToken, 6))))
	router.POST("/echorequest", rawHandler(responseHandler(definedChallengeHandler(receiveToken, 7))))
	router.POST("/testdisk", rawHandler(responseHandler(definedChallengeHandler(receiveToken, 8))))
	router.POST("/wg-step", rawHandler(responseHandler(definedChallengeHandler(receiveToken, 10))))
	router.POST("/api/recv-mail", remoteValidatorHandler(apiHandler(challengeAPIEMail)))
}

type givenToken struct {
	Login     string `json:"login"`
	Challenge int    `json:"challenge"`
	Token     string `json:"token"`
	token     []byte
	Data      []string `json:"data"`
}

func getChallengeList(_ httprouter.Params, _ []byte) (interface{}, error) {
	var ret []int
	for i := range challenges {
		ret = append(ret, i)
	}
	return ret, nil
}

func accessibleChallenge(r *http.Request, ps httprouter.Params, _ []byte) (interface{}, error) {
	if chid, err := strconv.Atoi(string(ps.ByName("chid"))); err != nil {
		return nil, err
	} else if chid == 0 || chid > len(challenges) {
		return nil, errors.New("This challenge doesn't exist")
	} else {
		for _, a := range challenges[chid].Accessible {
			if err := a(nil, r); err != nil {
				return nil, err
			}
		}
		return true, nil
	}
}

// receiveChallenge treats TP1{shadow,ssh_access} and TP2, TP3 challenges.
func receiveChallenge(r *http.Request, ps httprouter.Params, body []byte) (interface{}, error) {
	if chid, err := strconv.Atoi(string(ps.ByName("chid"))); err != nil {
		return nil, err
	} else if chid < 0 {
		return nil, errors.New("This challenge doesn't exist")
	} else if chid != 0 && chid < len(challenges) {
		return nil, errors.New("This is not the good way to hit this challenge")
	} else {
		var gt givenToken
		if err := json.Unmarshal(body, &gt); err != nil {
			return nil, err
		}

		if gt.Token != PongSecret {
			return nil, errors.New("This is not the expected token.")
		}

		var std *adlin.Student

		if stdid, err := strconv.Atoi(gt.Login); err == nil {
			if std, err = adlin.GetStudent(stdid); err != nil {
				return nil, err
			}
		} else if std, err = adlin.GetStudentByLogin(gt.Login); err != nil {
			return nil, err
		}

		if _, err := std.UnlockNewChallenge(chid, strings.Join(gt.Data, " ")); err != nil {
			if _, err := std.UpdateUnlockedChallenge(chid, strings.Join(gt.Data, " ")); err != nil {
				return nil, err
			}
		}

		return true, nil
	}
}

// receiveToken handles challenges for TP1 except shadow
func receiveToken(r *http.Request, body []byte, chid int) (interface{}, error) {
	var gt givenToken
	if err := json.Unmarshal(body, &gt); err != nil {
		return nil, err
	}

	gt.token = make([]byte, hex.DecodedLen(len(gt.Token)))
	if _, err := hex.Decode(gt.token, []byte(gt.Token)); err != nil {
		return nil, err
	}

	// Find challenge ID
	if gt.Challenge >= 1 {
		chid = gt.Challenge
	}

	if chid <= 0 || chid > len(challenges) {
		return nil, errors.New("This challenge doesn't exist")
	}

	// Is the challenge accessible?
	for _, a := range challenges[chid].Accessible {
		if err := a(nil, r); err != nil {
			return nil, err
		}
	}

	if std, err := adlin.GetStudentByLogin(gt.Login); err != nil {
		return nil, err
	} else {
		if err := challenges[chid].Check(std, &gt, chid); err != nil {
			log.Printf("%s just try ch#%d: %s\n", std.Login, chid, err)
			return nil, err
		}

		if _, err := std.UnlockNewChallenge(chid, gt.Token); err != nil {
			if me, ok := err.(*mysql.MySQLError); ok && me.Number == 1062 {
				return "Already validated", nil
			}
			log.Println(err)
			return nil, err
		}

		log.Printf("%s just unlock ch#%d\n", std.Login, chid)
		return "Success", nil
	}
}