package main

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

	"github.com/julienschmidt/httprouter"
)

const IPgwDMZ = "172.23.200.1"

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

/* Restrictions */

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

func accessFrom(ip string) func(_ *Student, r *http.Request) error {
	return func(_ *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(_ *Student, r *http.Request) error {
	return func(_ *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(_ *Student, r *http.Request) error {
	return func(_ *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(_ *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 challenge42(s *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 *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.")
	} else {
		return nil
	}
}

func challengeTime(s *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 *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 *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
	}

	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
	}
}

var challenges []Challenge

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

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

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

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

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

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

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

		/* Last : SSH key, see ssh.go:156 in NewKey function */
	}

	router.GET("/challenge", apiHandler(getChallengeList))
	router.GET("/challenge/:chid", rawHandler(accessibleChallenge))
	router.POST("/challenge", rawHandler(challengeHandler(receiveToken)))
	router.POST("/challenge/:chid", rawHandler(receiveChallenge))
	router.POST("/echorequest", rawHandler(definedChallengeHandler(receiveToken, 6)))
	router.POST("/testdisk", rawHandler(definedChallengeHandler(receiveToken, 7)))
}

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 - 1].Accessible {
			if err := a(nil, r); err != nil {
				return nil, err
			}
		}
		return true, nil
	}
}

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 <= 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 Student

		if stdid, err := strconv.Atoi(gt.Login); err == nil {
			if std, err = getStudent(stdid); err != nil {
				return nil, err
			}
		} else if std, err = 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
	}
}

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 - 1 > len(challenges) {
		return nil, errors.New("This challenge doesn't exist")
	}

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

	if std, err := getStudentByLogin(gt.Login); err != nil {
		return nil, err
	} else {
		if err := challenges[chid - 1].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 {
			log.Println(err)
			return nil, err
		}

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