package main

import (
	"crypto/rand"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os/exec"
	"strings"
	"time"

	"github.com/julienschmidt/httprouter"
)

func init() {
	router.GET("/api/wg.conf", func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
		w.Header().Set("Content-Type", "text/plain")

		err := GenWGConfig(w)

		if err != nil {
			w.Write([]byte(err.Error()))
		}
	})
	router.GET("/api/wg/", apiAuthHandler(showWgTunnel))
	router.GET("/api/wginfo", apiAuthHandler(func (student Student, ps httprouter.Params, body []byte) (interface{}, error) {
		return getTunnelInfo(student.Id), nil
	}))
	router.POST("/api/wg/", apiAuthHandler(genWgToken))
	router.POST("/api/wg/:token", getWgTunnelInfo)
}

func showWgTunnel(student Student, ps httprouter.Params, body []byte) (interface{}, error) {
	// Get tunnels assigned to the student
	return student.GetTunnelTokens()
}

func genWgToken(student Student, ps httprouter.Params, body []byte) (interface{}, error) {
	// Generate a token to access related wg info
	return student.NewTunnelToken()
}

type TunnelInfo struct {
	Status    string `json:"status"`
	SrvPubKey []byte `json:"srv_pubkey"`
	SrvPort   uint16 `json:"srv_port"`
	CltIPv6   net.IP `json:"clt_ipv6"`
	CltRange  uint8  `json:"clt_range"`
	SrvGW6    string `json:"srv_gw6"`
}

func getTunnelInfo(student int64) TunnelInfo {
	srv_pubkey, _ := base64.StdEncoding.DecodeString("uSpqyYovvP4OG6wDxZ0Qkq45MfyK58PMUuPaLesY8FI=")
	return TunnelInfo{
		Status: "OK",
		SrvPubKey: srv_pubkey,
		SrvPort: 42912,
		CltIPv6: studentIP(student),
		CltRange: 80,
		SrvGW6: "2a01:e0a:2b:2252::1",
	}
}

type PubTunnel struct {
	PubKey []byte
}

func getWgTunnelInfo(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	if addr := r.Header.Get("X-Forwarded-For"); addr != "" {
		r.RemoteAddr = addr
	}
	log.Printf("%s \"%s %s\" [%s]\n", r.RemoteAddr, r.Method, r.URL.Path, r.UserAgent())

	// Read the body
	if r.ContentLength < 0 || r.ContentLength > 6553600 {
		http.Error(w, "{errmsg:\"Request too large or request size unknown\"}", http.StatusRequestEntityTooLarge)
		return
	}

	// Access wg infos
	tokenhex := []byte(ps.ByName("token"))
	tokendec := make([]byte, hex.DecodedLen(len(tokenhex)))
	n, err := hex.Decode(tokendec, tokenhex)
	if err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusBadRequest)
		return
	}

	token, err := GetTunnelToken(tokendec[:n])
	if err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusBadRequest)
		return
	}

	var pt PubTunnel
	if err := json.NewDecoder(r.Body).Decode(&pt); err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusBadRequest)
		return
	}

	token.PubKey = pt.PubKey
	_, err = token.Update()
	if err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusBadRequest)
		return
	}

	syncWgConf()

	tinfo := getTunnelInfo(token.IdStudent)

	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte(fmt.Sprintf(`[Peer]
PublicKey = %s
Endpoint = %s:%d
AllowedIPs = ::/0
PersistentKeepalive = 5
# MyIPv6=%s1/%d
# MyNetwork=%s/%d
# GWIPv6=%s
`, base64.StdEncoding.EncodeToString(tinfo.SrvPubKey), "82.64.31.248", tinfo.SrvPort, tinfo.CltIPv6, 64, tinfo.CltIPv6, tinfo.CltRange, tinfo.SrvGW6)))
}


type TunnelToken struct {
	token     []byte
	TokenText string
	IdStudent int64
	PubKey    []byte
	Time      time.Time
	Dump      *WGDump
}

func GetTunnelToken(token []byte) (t TunnelToken, err error) {
	err = DBQueryRow("SELECT token, token_text, id_student, pubkey, time FROM student_tunnel_tokens WHERE token=? ORDER BY time DESC", token).Scan(&t.token, &t.TokenText, &t.IdStudent, &t.PubKey, &t.Time)
	if err == nil && t.PubKey != nil {
		if wgd, errr := readWgDump(); errr == nil {
			if v, ok := wgd[base64.StdEncoding.EncodeToString(t.PubKey)]; ok {
				t.Dump = &v
			}
		}
	}
	return
}

func tokenFromText(token string) []byte {
	sha := sha512.Sum512([]byte(token))
	return sha[:]
}

func (student Student) NewTunnelToken() (t TunnelToken, err error) {
	tok := make([]byte, 7)
	if _, err = rand.Read(tok); err != nil {
		return
	}

	t.TokenText = strings.Replace(strings.Replace(strings.Replace(strings.Replace(strings.Replace(base64.RawStdEncoding.EncodeToString(tok), "/", ".", -1), "+", "_", -1), "O", "#", -1), "l", "$", -1), "I", ">", -1)
	t.token = tokenFromText(t.TokenText)
	t.IdStudent = student.Id

	_, err = DBExec("INSERT INTO student_tunnel_tokens (token, token_text, id_student, time) VALUES (?, ?, ?, ?)", t.token, t.TokenText, student.Id, time.Now())
	return
}

func (student Student) GetTunnelTokens() (ts []TunnelToken, err error) {
	if rows, errr := DBQuery("SELECT token, token_text, id_student, pubkey, time FROM student_tunnel_tokens WHERE id_student = ? ORDER BY time DESC", student.Id); errr != nil {
		return nil, errr
	} else if wgd, errr := readWgDump(); errr != nil {
		return nil, errr
	} else {
		defer rows.Close()

		for rows.Next() {
			var t TunnelToken
			if err = rows.Scan(&t.token, &t.TokenText, &t.IdStudent, &t.PubKey, &t.Time); err != nil {
				return
			}
			if t.PubKey != nil {
				if v, ok := wgd[base64.StdEncoding.EncodeToString(t.PubKey)]; ok {
					t.Dump = &v
				}
			}
			ts = append(ts, t)
		}
		if err = rows.Err(); err != nil {
			return
		}

		return
	}
}

func (student Student) GetTunnelToken(token []byte) (t TunnelToken, err error) {
	err = DBQueryRow("SELECT token, token_text, id_student, pubkey, time FROM student_tunnel_tokens WHERE token = ? AND id_student = ? ORDER BY time DESC", token, student.Id).Scan(&t.token, &t.TokenText, &t.IdStudent, &t.PubKey, &t.Time)
	if err == nil && t.PubKey != nil {
		if wgd, errr := readWgDump(); errr == nil {
			if v, ok := wgd[base64.StdEncoding.EncodeToString(t.PubKey)]; ok {
				t.Dump = &v
			}
		}
	}
	return
}

func (t *TunnelToken) Update() (int64, error) {
	newtoken := tokenFromText(t.TokenText)
	tm := time.Now()

	if res, err := DBExec("UPDATE student_tunnel_tokens SET token = ?, token_text = ?, id_student = ?, pubkey = ?, time = ? WHERE token = ?", newtoken, t.TokenText, t.IdStudent, t.PubKey, tm, t.token); err != nil {
		return 0, err
	} else if nb, err := res.RowsAffected(); err != nil {
		return 0, err
	} else {
		t.token = newtoken
		t.Time = tm
		return nb, err
	}
}

func GetStudentsTunnels() (ts []TunnelToken, err error) {
	if rows, errr := DBQuery("SELECT T.token, T.token_text, T.id_student, T.pubkey, T.time FROM student_tunnel_tokens T INNER JOIN (SELECT B.id_student, MAX(B.time) AS time FROM student_tunnel_tokens B WHERE B.pubkey IS NOT NULL GROUP BY id_student) L ON T.id_student = L.id_student AND T.time = L.time"); errr != nil {
		return nil, errr
	} else {
		defer rows.Close()

		for rows.Next() {
			var t TunnelToken
			if err = rows.Scan(&t.token, &t.TokenText, &t.IdStudent, &t.PubKey, &t.Time); err != nil {
				return
			}
			ts = append(ts, t)
		}

		err = rows.Err()
		return
	}
}

func studentIP(idstd int64) net.IP {
	return net.ParseIP(fmt.Sprintf("2a01:e0a:2b:2252:%x::", idstd))
}

func GenWGConfig(w io.Writer) (error) {
	ts, err := GetStudentsTunnels()
	if err != nil {
		return err
	}

	for _, t := range ts {
		if t.PubKey == nil {
			continue
		}

		w.Write([]byte(fmt.Sprintf(`[Peer]
#IdStudent = %d
PublicKey = %s
AllowedIPs = %s/%d
`, t.IdStudent, base64.StdEncoding.EncodeToString(t.PubKey), studentIP(t.IdStudent), 80)))
	}

	return nil
}


func syncWgConf() (err error) {
	_, err = exec.Command("sh", "/root/wg-sync.sh").Output()
	return
}


type WGDump struct {
	PubKey     string
	PSK        string
	Endpoint   string
	AllowedIPs string
	LastHandS  string
	RX         string
	TX         string
	KeepAlive  string
}

func readWgDump() (wgd map[string]WGDump, err error) {
	out, errr := exec.Command("wg", "show", "wg-adlin", "dump").Output()

	if errr != nil {
		return nil, errr
	}

	wgd = map[string]WGDump{}
	for _, line := range strings.Split(string(out), "\n") {
		cols := strings.Fields(line)
		if len(cols) != 8 {
			continue
		}

		wgd[cols[0]] = WGDump{cols[0], cols[1], cols[2], cols[3], cols[4], cols[5], cols[6], cols[7]}
	}

	return
}