package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	_ "github.com/dgrijalva/jwt-go"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

const (
	NONCE_SIZE        int    = 16
	TOKEN_COOKIE_NAME string = "sso_token"
	RFP_COOKIE_NAME   string = "sso_rfp"
)

// acr_values can be mfa or selective_mfa (mfa only for external users)
// mfa amr values:
// pas - password
// otp - OTP code
// u2f - U2F code
// mfa - multi-factor
// hrd - hardware OTP device used
// sft - software OTP device used

type ProxyConfig struct {
	IDProviderURL      string
	ClientID           string
	UpstreamURL        string
	ListenOn           string
	TrustedCACert      string
	RequestMFA         bool
	AllowedMFAMethods  []string // An OR set
	RequiredMFAMethods []string // An AND set
	reverseProxy       *httputil.ReverseProxy
}

func URLMustParse(u string) *url.URL {
	o, err := url.Parse(u)
	if err != nil {
		panic(err)
	}
	return o
}

func GenerateNonce() (string, error) {
	nonce := make([]byte, NONCE_SIZE)
	n, err := rand.Read(nonce)
	if n != NONCE_SIZE || err != nil {
		return "", err
	}
	return hex.EncodeToString(nonce), nil
}

func CompareUpper(lhs, rhs string) bool {
	return strings.ToUpper(lhs) == strings.ToUpper(rhs)
}

// TODO
// Cookie rules
// Secure
// HttpOnly
// Path to /
// Expires to iat in JWT
func SetCookie() {
}

// TODO
// Fetch (connect timeout 1s, read timeout 30s, read size 1M)
func DownloadCertificate() {
}

// TODO
// Fetch (connect timeout 1s, read timeout 30s, read size 1M)
func DownloadCRL() {
}

// TODO
// Cert validation
// Validate cert not in CRL
// Validate cert chains to trusted CA cert (ship with proxy)
// Validate CRL signed by trusted CA
// Cert "Subject CN " must match exactly PKI setting (ex: foo-pki.foo.com)
// Current time bust be within "Validity Not Before" and "Validity Not After" in cert +- 5 minutes
// Cert Key length >= 2048
// Certificate usage must include "digitalSignature"
func ValidateCertificate() {
}

func MakeClientID(r *http.Request) string {
	if strings.Contains(r.Host, ":") {
		return r.Host
	}
}

// TODO
func RedirectToIDP(w http.ResponseWriter, r *http.Request) {
	nonce := GenerateNonce()
	nonceh := "" // SHA256 nonce

	// Set nonce cookie

	req := url.Values{}
	req.Add("client_id", "") // fqdn + : + port
	req.Add("nonce", nonceh)
	req.Add("redirect_uri", "") // Requested URL
	req.Add("scope", "openid")
	req.Add("response_type", "id_token")
}

// TODO: Remove id_token from URL, set cookie and redirect user to requested URL
func SetTokenCookieAndRedirect(w http.ResponseWriter, r *http.Request, token string) {
}

// TODO
// Fetch cert from x5u URL
// Get CRL from cert, fetch (connect timeout 1s, read timeout 30s, read size 1M)
//
// exp claim has passed +- 5 minutes
// iat claim is greater than 24 hours +- 5 minutes
// aud claim is exact match for client_id
// iss claim is exact match for idp (ex: foo.example.com
// if other aud claims validate that they are known
// nonce in JWT must be SHA256 of rfp cookie value
// Validate cert
// alg jwt header must be one of [PS256, PS385, PS512]
// typ jwt header must be JWS
// validate jwt signature
// validate amr claim contains requested acr values (selective_mfa will be just mfa)
// validate acr claim is the same as requested acr_values
func ValidateJWT(jwt, rfp string) bool {
	return true
}

// TODO
func GetJWTSubject(jwt string) string {
	return ""
}

func RequestHasForwardedUser(w http.ResponseWriter, r *http.Request) bool {
	if _, ok := r.Header["X-Forwarded-User"]; ok {
		log.Printf("ERROR: Request contains X-Forwarded-For header")
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return true
	} else {
		return false
	}
}

func RequestIsOverSecureChannel(w http.ResponseWriter, r *http.Request) {
	if https, ok := r.Header["X-Forwarded-Proto"]; !ok || len(https) != 1 {
		log.Printf("ERROR: Request does not contain X-Forwarded-Proto header")
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return false
	}

	if !CompareUpper(https[0], "HTTPS") {
		log.Printf("ERROR: Request is not over HTTPS")
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return false
	}

	return true
}

// TODO
// - Validate Hostname header == known hostname
func AuthProxyController(w http.ResponseWriter, r *http.Request) {
	proxy := r.Context().Value("ProxyConfig").(*ProxyConfig).reverseProxy

	// Order matters in these checks!
	if RequestHasForwardedUser(w, r) {
		return
	}

	if !RequestIsOverSecureChannel(w, r) {
		return
	}

	if CompareUpper(r.Method, "OPTIONS") {
		proxy.ServeHTTP(w, r)
		return
	}

	rfpc, err := r.Cookie(RFP_COOKIE_NAME)
	if err != nil {
		log.Printf("ERROR: No rfp cookie")
		RedirectToIDP(w, r)
		return
	}

	token := r.URL.Query().Get("id_token")
	if token != "" && ValidateJWT(rfpc.Value, token) {
		SetTokenCookieAndRedirect(w, r, token)
		return
	}

	tokenc, err := r.Cookie(TOKEN_COOKIE_NAME)
	if err != nil {
		log.Printf("ERROR: No token cookie")
		RedirectToIDP()
		return
	}

	if !ValidateJWT(tokenc.Value, rfpc.Value) {
		log.Printf("ERROR: Token is invalid")
		RedirectToIDP()
		return
	}

	r.Header["X-Forwarded-User"] = []string{GetJWTSubject(tokenc.Value)}

	proxy.ServeHTTP(w, r)
}

// Remove token and rfp cookies and redirect user to root of domain
func LogoutController(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:   TOKEN_COOKIE_NAME,
		Value:  "",
		MaxAge: 0,
	})

	http.SetCookie(w, &http.Cookie{
		Name:   TOKEN_COOKIE_NAME,
		Value:  "",
		MaxAge: 0,
	})

	http.Redirect(w, r, "/", http.StatusFound)
}

// TODO
func parseConfig() *ProxyConfig {
	return &ProxyConfig{
		IDProviderURL: "",
		ClientID:      "",
		UpstreamURL:   "http://localhost:9991/",
		ListenOn:      ":9992",
		TrustedCACert: "",
	}
}

func main() {
	cfg := parseConfig()
	cfg.reverseProxy = httputil.NewSingleHostReverseProxy(URLMustParse(cfg.UpstreamURL))

	http.HandleFunc("/.oidc/logout", func(w http.ResponseWriter, r *http.Request) {
		LogoutController(w,
			r.WithContext(context.WithValue(r.Context(), "ProxyConfig", cfg)))
	})

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		AuthProxyController(w,
			r.WithContext(context.WithValue(r.Context(), "ProxyConfig", cfg)))
	})

	log.Fatal(http.ListenAndServe(cfg.ListenOn, nil))
}