package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"gopkg.in/square/go-jose.v2"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"
)

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

// TODO: Enable https checks in HTTP client

// 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
	PKISubject         string // TODO: Should be same as IDP w/out scheme and port
	ClockSkew          time.Duration
	MaxLiftetime       time.Duration
	IsOptional         bool
	RequestMFA         bool
	AllowedMFAMethods  []string // An OR set
	RequiredMFAMethods []string // An AND set
	reverseProxy       *httputil.ReverseProxy
}

type IdPConfig struct {
	AuthorizationEndpoint string   `json:"authorization_endpoint"`
	Issuer                string   `json:"issuer"`
	JwksUri               string   `json:"jwks_uri"`
	GrantTypes            []string `json:"grant_types_supported"`
	IdTokenSigningAlgs    []string `json:"id_token_signing_alg_values_supported"`
	ResponseModes         []string `json:"response_modes_supported"`
	ResponseTypes         []string `json:"response_types_supported"`
	Scopes                []string `json:"scopes_supported"`
	SubjectTypes          []string `json:"subject_types_supported"`
}

// TODO: Optimization to fetch only if expired (per http headers)
func FetchIdPConfig(h CautiousHTTPClient, idp_url string) (*IdPConfig, error) {
	u, err := url.Parse(idp_url)
	if err != nil {
		return nil, err
	}
	u.Path = "/.well-known/openid-configuration"

	var idpc IdPConfig
	err = h.GetJSON(u.String(), &idpc)
	if err != nil {
		return nil, err
	}

	return &idpc, nil
}

// TODO: Optimization to fetch only if expired (per http headers)
func FetchJWKS(h CautiousHTTPClient, jwks_url string, val KeyValidator) (map[string]jose.JSONWebKey, error) {
	var jwks jose.JSONWebKeySet
	err := h.GetJSON(jwks_url, &jwks)
	if err != nil {
		return nil, err
	}

	keys := make(map[string]jose.JSONWebKey, len(jwks.Keys))

	for _, k := range jwks.Keys {
		err = val.Validate(k)
		if err == nil {
			keys[k.KeyID] = k
		}
	}

	return keys, nil
}

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
}

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

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

// TODO
func RedirectToIDP(w http.ResponseWriter, r *http.Request) {
	nonce, _ := GenerateNonce()
	_ = nonce
	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
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) bool {
	https, ok := r.Header["X-Forwarded-Proto"]
	if !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(w, r)
		return
	}

	if !ValidateJWT(tokenc.Value, rfpc.Value) {
		log.Printf("ERROR: Token is invalid")
		RedirectToIDP(w, r)
		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 LoginController(w http.ResponseWriter, r *http.Request) {
}

// TODO
// Optional login allows for applications that can operate in anonymous mode or
// authenticated mode. When in anonmyous mode the request is proxied through
// without an X-Forwarded-User header. Upstream servers should either expose or
// map a URL for /.oidc/login to allow users to login. On successful login the
// user will be redirected back to the main page for the site (/)
func parseConfig() *ProxyConfig {
	return &ProxyConfig{
		IDProviderURL: "http://mcrute-virt:9993",
		ClientID:      "test.crute.me:443",
		UpstreamURL:   "http://localhost:9991/",
		ListenOn:      ":9992",
		TrustedCACert: "/home/mcrute/oidc_project/test_ca/ca_cert.pem",
		IsOptional:    false,
		PKISubject:    "Crute OpenID Signing 1",
		MaxLiftetime:  24 * time.Hour,
		ClockSkew:     5 * time.Minute,
	}
}

func main() {
	cfg := parseConfig()
	h := NewCautiousHTTPClient()

	v := NewKeyValidator(cfg.PKISubject)
	v.LoadRootPEM(cfg.TrustedCACert)

	idpc, err := FetchIdPConfig(h, cfg.IDProviderURL)
	if err != nil {
		fmt.Printf("%s\n", err)
		return
	}

	jwks, err := FetchJWKS(h, idpc.JwksUri, v)
	if err != nil {
		fmt.Printf("%s\n", err)
		return
	}

	jv := NewJWSValidator(jwks, idpc.Issuer, cfg.ClientID, cfg.ClockSkew, cfg.MaxLiftetime)

	nonce := "ofspmfjuvoswhhde"
	raw_jwt := "eyJ0eXAiOiJKV1MiLCJhbGciOiJQUzI1NiIsImtpZCI6IjEifQ.eyJub25jZSI6IjM0MjlhMjAyYzU4ZDkyYjQwNjNjOWM4MWM2MjQyNGRlNzBkMmIzZDQ4MmVlNDFhOTdjYmNhZjEwZDk5MWFiOTMiLCJpc3MiOiJpZHAuY3J1dGUubWU6NDQzIiwiaWF0IjoxNTA0NTc2Mzc0LCJuYmYiOjE1MDQ1NzYzNzQsImV4cCI6MTUwNDY2Mjc3NCwic3ViIjoibWNydXRlIiwiYXVkIjoidGVzdC5jcnV0ZS5tZTo0NDMifQ.iizlNfY1Vg7d-XRmgyYuhpNkNrOGaT9OOgO0HdjBozOWMvKzBTtATbIfoWOrNH6DiFY1as8uy3I1Pxnkrb8Ti8_cLDQeLxOv9klAbnebeuPI_wtZ0iwSUnSWaYzN6I6sqcEjHX3fibFvAQhO5dNDzSwONjw4AvcdpZKh579FO1sAvIw-1DmMyPSUun7rbC0Kf1Jtdlr3q7tOp3wdI_erkstxCNPwyuv7X1J7uetsu0BeJS25C2DxeB03BPEIUoo_C1xvcqikfSLLpoFcyToYiS-R9o-WpRjGid_yug65J5ALn2aM3vhe9rRbydKVm_omGL8-Etj06zbqM0Y6OrJUgA"
	claims, err := jv.Validate(raw_jwt, nonce)
	if err != nil {
		fmt.Printf("Error validating: %s\n", err)
		return
	}

	fmt.Printf("Valid JWT for: %+v\n", claims.Subject)

	return

	cfg.reverseProxy = httputil.NewSingleHostReverseProxy(URLMustParse(cfg.UpstreamURL))

	if cfg.IsOptional {
		http.HandleFunc("/.oidc/login", func(w http.ResponseWriter, r *http.Request) {
			LoginController(w,
				r.WithContext(context.WithValue(r.Context(), "ProxyConfig", cfg)))
		})
	}

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