package main

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base32"
	"encoding/binary"
	"encoding/hex"
	"encoding/pem"
	"flag"
	"fmt"
	"github.com/ipfs/go-cid"
	"github.com/mroth/weightedrand"
	"github.com/multiformats/go-multihash"
	"log"
	"math/big"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

var base32Encoding = base32.NewEncoding("qpzry9x8gf2tvdw0s3jn54khce6mua7l").WithPadding(base32.NoPadding)

var sniAddress string

var privateKey ed25519.PrivateKey
var messageCacheLimit = 4096
var signedMessageCacheLimit int
var messageCacheMutex sync.RWMutex
var messageCache = make(map[string]*ContentMessage)

var contentServers []*ContentServer

type ContentServer struct {
	Index           int
	Address         string
	Weight          uint
	LastCheckResult bool
	LastCheck       time.Time
	LastCheckMutex  sync.RWMutex
}

func (s *ContentServer) getURL(args ...string) string {
	return fmt.Sprintf("https://%s/%s", s.Address, strings.Join(args, "/"))
}

func (s *ContentServer) getCheckResult() bool {
	s.LastCheckMutex.RLock()
	defer s.LastCheckMutex.RUnlock()
	return s.LastCheckResult
}

func (s *ContentServer) setCheckResult(result bool) {
	s.LastCheckMutex.Lock()
	defer s.LastCheckMutex.Unlock()
	s.LastCheckResult = result
}

func (s *ContentServer) check() {
	customTransport := http.DefaultTransport.(*http.Transport).Clone()
	customTransport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	client := &http.Client{
		Transport: customTransport,
		Timeout:   5 * time.Second,
	}
	response, err := client.Head(s.getURL())

	if err != nil {
		s.setCheckResult(false)
		return
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusBadRequest {
		s.setCheckResult(false)
		return
	}
	s.setCheckResult(true)
}

func selectNextContentServer(skip []int) *ContentServer {
	inSkip := func(i int) bool {
		for _, c := range skip {
			if i == c {
				return true
			}
		}
		return false
	}

	var choosingList []weightedrand.Choice
	for _, c := range contentServers {
		if !inSkip(c.Index) && c.getCheckResult() {
			choosingList = append(choosingList, weightedrand.NewChoice(c, c.Weight))
		}
	}

	chooser, err := weightedrand.NewChooser(choosingList...)

	if err != nil {
		return nil
	}

	return chooser.Pick().(*ContentServer)
}

//TODO: move this to a library
type ContentMessage struct {
	Version            uint64
	PublicKey          ed25519.PublicKey
	IssueTime          int64
	Identifier         cid.Cid
	VerificationResult *bool
	Signature          []byte
}

func (s *ContentMessage) sign(privateKey ed25519.PrivateKey) {
	s.PublicKey = make([]byte, ed25519.PublicKeySize)
	copy(s.PublicKey, privateKey[32:])
	s.Signature = ed25519.Sign(privateKey, s.encodeMessage())
	s.VerificationResult = nil
}

func (s *ContentMessage) verify() bool {
	currentTime := time.Now()

	notBefore := currentTime.Add(-time.Hour)    //Only one hour before time
	notAfter := currentTime.Add(time.Hour * 24) //Only 24 hours after time
	issueTime := time.Unix(s.IssueTime, 0)

	if issueTime.Before(notBefore) {
		return false
	}

	if issueTime.After(notAfter) {
		return false
	}

	messageCacheMutex.RLock()
	k := string(s.encode())
	cachedMessage, ok := messageCache[k]
	messageCacheMutex.RUnlock()
	if ok {
		return *cachedMessage.VerificationResult
	}

	messageCacheMutex.Lock()
	defer messageCacheMutex.Unlock()

	if s.VerificationResult == nil {
		makeBool := func(v bool) *bool { return &v }
		s.VerificationResult = makeBool(ed25519.Verify(s.PublicKey, s.encodeMessage(), s.Signature))
	}

	if len(messageCache) >= messageCacheLimit {
		//Find oldest value, remove it
		var item *ContentMessage
		for _, e := range messageCache {
			if item == nil || e.IssueTime < item.IssueTime {
				item = e
			}
		}

		if item != nil {
			delete(messageCache, string(item.encode()))
		}
	}

	messageCache[k] = s

	return *s.VerificationResult
}

func (s *ContentMessage) encodeMessage() []byte {
	message := &bytes.Buffer{}

	buf := make([]byte, binary.MaxVarintLen64)

	n := binary.PutUvarint(buf, s.Version) //signature version
	_, _ = message.Write(buf[:n])

	if s.Version == 0 {
		_, _ = message.Write(s.PublicKey)
		n = binary.PutVarint(buf, s.IssueTime) //time
		_, _ = message.Write(buf[:n])

		_, _ = s.Identifier.WriteBytes(message)

		return message.Bytes()
	}

	return nil
}

func (s *ContentMessage) encode() []byte {
	message := s.encodeMessage()

	if message == nil || len(s.Signature) != ed25519.SignatureSize {
		return nil
	}

	return append(message, s.Signature...)
}

func (s ContentMessage) String() string {
	return fmt.Sprintf("%d %x %d %s %x", s.Version, s.PublicKey, s.IssueTime, s.Identifier.String(), s.Signature)
}

func DecodeContentMessage(signatureBytes []byte) *ContentMessage {
	message := ContentMessage{}

	buffer := bytes.NewBuffer(signatureBytes)

	var err error

	message.Version, err = binary.ReadUvarint(buffer)
	if err != nil {
		log.Print(err)
		return nil
	}

	if message.Version == 0 {
		message.PublicKey = make([]byte, ed25519.PublicKeySize)
		_, err := buffer.Read(message.PublicKey)
		if err != nil {
			log.Print(message.String())
			log.Print(err)
			return nil
		}
		message.IssueTime, err = binary.ReadVarint(buffer)
		if err != nil {
			log.Print(message.String())
			log.Print(err)
			return nil
		}

		_, message.Identifier, err = cid.CidFromReader(buffer)
		if err != nil {
			log.Print(message.String())
			log.Print(err)
			return nil
		}

		message.Signature = make([]byte, ed25519.SignatureSize)
		_, err = buffer.Read(message.Signature)
		if err != nil {
			log.Print(message.String())
			log.Print(err)
			return nil
		}

		if buffer.Len() != 0 { //Unknown extra data
			log.Print(message.String())
			log.Printf("%x", buffer.Bytes())
			return nil
		}

		return &message

	}

	return nil
}

func createSelfSignedCertificate() ([]byte, []byte) {
	x509Template := x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{},
		NotBefore:             time.Unix(0, 0).UTC(),
		NotAfter:              time.Date(time.Now().UTC().Year()+10, 0, 0, 0, 0, 0, 0, time.UTC),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLen:            1,
	}

	privateBogusKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	if err != nil {
		log.Fatal(err)
	}

	certBytes, err := x509.CreateCertificate(rand.Reader, &x509Template, &x509Template, privateBogusKey.Public(), privateBogusKey)
	if err != nil {
		log.Fatal(err)
	}

	keyBytes, err := x509.MarshalPKCS8PrivateKey(privateBogusKey)
	if err != nil {
		log.Fatal(err)
	}

	return pem.EncodeToMemory(&pem.Block{
			Type:  "CERTIFICATE",
			Bytes: certBytes,
		}),
		pem.EncodeToMemory(&pem.Block{
			Type:  "PRIVATE KEY",
			Bytes: keyBytes,
		})
}

func isSNIAllowed(sni string) bool {
	return len(sniAddress) == 0 || sni == sniAddress
}

func setOtherHeaders(w http.ResponseWriter) {
	w.Header().Set("Server", "FinalCommander")
	w.Header().Set("Vary", "Content-Encoding")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-Robots-Tags", "noindex, nofollow, notranslate")
	w.Header().Set("Referrer-Policy", "origin")
}
func setCORSHeaders(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Credentials", "true")
	w.Header().Set("Access-Control-Max-Age", "7200") //Firefox caps this to 86400, Chrome to 7200. Default is 5 seconds (!!!)
	w.Header().Set("Access-Control-Allow-Methods", "GET,HEAD,OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "DNT,ETag,Origin,Accept,Accept-Language,X-Requested-With,Range")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Expose-Headers", "*")

	//CORP, COEP, COOP
	w.Header().Set("Cross-Origin-Embedder-Policy", "require-corp")
	w.Header().Set("Cross-Origin-Resource-Policy", "cross-origin")
	w.Header().Set("Cross-Origin-Opener-Policy", "unsafe-none")
}

func encodeIntegerList(data []int) []byte {
	message := &bytes.Buffer{}
	buf := make([]byte, binary.MaxVarintLen64)

	for _, i := range data {
		n := binary.PutVarint(buf, int64(i))
		_, _ = message.Write(buf[:n])
	}

	return message.Bytes()
}
func decodeIntegerList(data []byte) []int {
	buf := bytes.NewBuffer(data)
	var result []int

	for {
		if buf.Len() <= 0 {
			break
		}

		i, err := binary.ReadVarint(buf)
		if err != nil {
			//TODO: maybe should error
			break
		}
		result = append(result, int(i))
	}

	return result
}

func handle(w http.ResponseWriter, r *http.Request) {
	if len(r.Host) > 0 && r.Host == r.TLS.ServerName { //Prevents rebinding / DNS stuff
		w.WriteHeader(http.StatusNotFound)
		return
	}

	if r.Method == "GET" || r.Method == "HEAD" {
		log.Printf("Serve %s", r.URL.Path)
		setOtherHeaders(w)
		setCORSHeaders(w)
		pathElements := strings.Split(r.URL.Path, "/")
		if len(pathElements) < 3 {
			log.Printf("1")
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		hashType := strings.ToLower(pathElements[1])

		hash, err := hex.DecodeString(pathElements[2])
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		var skip []int

		if len(pathElements) > 3 {
			data, err := base32Encoding.DecodeString(pathElements[3])
			if err != nil {
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			skip = decodeIntegerList(data)
		}

		var identifier cid.Cid

		if hashType == "sha256" && len(hash) == 32 {
			mh, _ := multihash.Encode(hash, multihash.SHA2_256)
			identifier = cid.NewCidV1(cid.Raw, mh)
		} else if hashType == "md5" && len(hash) == 16 {
			mh, _ := multihash.Encode(hash, multihash.MD5)
			identifier = cid.NewCidV1(cid.Raw, mh)
		} else {
			w.WriteHeader(http.StatusNotFound)
			return
		}

		contentServer := selectNextContentServer(skip)
		if contentServer == nil {
			w.WriteHeader(http.StatusNotFound)
			return
		}

		//TODO: signing cache
		message := &ContentMessage{
			Version:    0,
			IssueTime:  time.Now().UTC().Unix(),
			Identifier: identifier,
		}
		message.sign(privateKey)

		skip = append(skip, contentServer.Index)

		http.Redirect(w, r, contentServer.getURL(base32Encoding.EncodeToString(message.encode()), base32Encoding.EncodeToString(encodeIntegerList(skip))), http.StatusFound)
	} else if r.Method == "OPTIONS" {
		setOtherHeaders(w)
		setCORSHeaders(w)
		w.WriteHeader(http.StatusNoContent)
	} else {
		w.WriteHeader(http.StatusNotImplemented)
	}
}

func checkContentServers() {
	checkTime := time.Now().Add(-15 * time.Minute)
	for _, c := range contentServers {
		if c.LastCheck.Before(checkTime) {
			c.check()
		}
	}
}

func main() {
	//TODO: OCSP
	certificatePath := flag.String("certificate", "ssl.crt", "Path to SSL certificate file.")
	keypairPath := flag.String("keypair", "ssl.key", "Path to SSL key file.")

	listenAddress := flag.String("listen", ":7777", "Address/port to lisent on.")

	weightedServerList := flag.String("servers", "", "Weighted list of servers to use. All use HTTPs. Format address:PORT/WEIGHT,[...]")

	sniAddressOption := flag.String("sni", "", "Define SNI address if desired. Empty will serve any requests regardless.")

	signatureCacheLimitOption := flag.Int("siglimit", 128, "Maximum number of lingering valid signature produced cached.")

	flag.Parse()

	signedMessageCacheLimit = *signatureCacheLimitOption

	var err error

	privateKeyEnv := os.Getenv("PRIVATE_KEY")

	if privateKeyEnv == "" {
		log.Print("No PRIVATE_KEY environment variable specified, generating new identity")
		publicKey, privateKey, _ := ed25519.GenerateKey(rand.Reader)

		log.Printf("Public Ed25519 key (share this): %s", base32Encoding.EncodeToString(publicKey))
		log.Printf("Private Ed25519 key (keep safe!): %s", base32Encoding.EncodeToString(privateKey))
		return
	}

	privateKey, err = base32Encoding.DecodeString(privateKeyEnv)
	if err != nil {
		log.Fatal(err)
	}
	publicKey := make([]byte, ed25519.PublicKeySize)
	copy(publicKey, privateKey[32:])
	log.Printf("Loaded Private Ed25519 key, Public %s", base32Encoding.EncodeToString(publicKey))

	if len(privateKey) != ed25519.PrivateKeySize {
		log.Fatal("Wrong Private key length")
	}

	for i, s := range strings.Split(*weightedServerList, ",") {
		p := strings.Split(s, "/")
		if len(p) != 2 {
			log.Fatalf("Invalid weighted server %s", s)
		}

		weight, err := strconv.ParseUint(p[1], 10, 32)
		if err != nil {
			log.Fatal(err)
		}

		contentServers = append(contentServers, &ContentServer{
			Index:           i,
			Address:         p[0],
			Weight:          uint(weight),
			LastCheckResult: false,
			LastCheck:       time.Date(1970, 0, 0, 0, 0, 0, 0, time.UTC),
		})
	}

	//TODO: cron this
	checkContentServers()

	sniAddress = strings.ToLower(*sniAddressOption)

	bogusCertificatePEM, bogusKeyPairPEM := createSelfSignedCertificate()

	bogusCertificate, err := tls.X509KeyPair(bogusCertificatePEM, bogusKeyPairPEM)
	if err != nil {
		log.Fatal(err)
	}
	serverCertificate, err := tls.LoadX509KeyPair(*certificatePath, *keypairPath)
	if err != nil {
		log.Fatal(err)
	}

	server := &http.Server{
		Addr:    *listenAddress,
		Handler: http.HandlerFunc(handle),
		TLSConfig: &tls.Config{
			MinVersion: tls.VersionTLS12,
			MaxVersion: 0, //max supported, currently TLS 1.3
			CurvePreferences: []tls.CurveID{
				tls.X25519,
				tls.CurveP256,
				tls.CurveP384,
			},
			CipherSuites: []uint16{
				tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
				tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
				tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			},
			PreferServerCipherSuites: false,
			SessionTicketsDisabled:   false,
			Renegotiation:            tls.RenegotiateFreelyAsClient,
			NextProtos: []string{
				"h2",
				"http/1.1",
			},
			GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
				if isSNIAllowed(info.ServerName) {
					return &serverCertificate, nil
				}
				return &bogusCertificate, nil
			},
		},
	}

	log.Printf("Serving on %s", *listenAddress)

	log.Fatal(server.ListenAndServeTLS(*certificatePath, *keypairPath))
}