package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"flag"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/gorilla/handlers"
	"github.com/pkg/errors"
)

const (
	NONCE_SIZE           = 16
	TOKEN_COOKIE_NAME    = "sso_token"
	RFP_COOKIE_NAME      = "sso_rfp"
	DEFAULT_CLOCK_SKEW   = 5 * time.Minute
	DEFAULT_MAX_LIFETIME = 24 * time.Hour
	DEFAULT_COOKIE_EXP   = 48 * time.Hour
)

// TODO: MFA support

type ProxyConfig struct {
	IdProviderURL          *url.URL
	IdProviderAuthEndpoint *url.URL
	ClientId               *url.URL
	UpstreamURL            string
	ListenOn               string
	TrustedCACert          string
	ClockSkew              time.Duration
	MaxLiftetime           time.Duration
	IsOptional             bool
	IsBootstrap            bool
	RequestMFA             bool
	AllowedMFAMethods      []string // An OR set
	RequiredMFAMethods     []string // An AND set
	reverseProxy           *httputil.ReverseProxy
	jwsValidator           JWSValidator
}

type IdPConfig struct {
	AuthorizationEndpoint *JSONURL `json:"authorization_endpoint"`
	JwksUri               *JSONURL `json:"jwks_uri"`
	Issuer                string   `json:"issuer"`
	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"`
}

func FetchIdPConfig(h CautiousHTTPClient, u *url.URL) (*IdPConfig, error) {
	u = URLMustParse(u.String())
	u.Path = "/.well-known/openid-configuration"

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

	return &idpc, nil
}

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

func SetSecureCookie(w http.ResponseWriter, name string, value string, exp time.Duration) {
	http.SetCookie(w, &http.Cookie{
		Name:     name,
		Value:    value,
		Expires:  time.Now().Add(exp),
		HttpOnly: true,
		Secure:   true,
		Path:     "/",
	})
}

func ExpireCookie(w http.ResponseWriter, name string) {
	http.SetCookie(w, &http.Cookie{
		Name:     name,
		Value:    "",
		Expires:  time.Now().Add(-1 * time.Hour),
		HttpOnly: true,
		Secure:   true,
		Path:     "/",
		MaxAge:   0,
	})
}

func RedirectToIdP(w http.ResponseWriter, r *http.Request, path string) {
	ctx := r.Context().Value("ProxyConfig").(*ProxyConfig)

	nonce, err := GenerateNonce()
	if err != nil {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	SetSecureCookie(w, RFP_COOKIE_NAME, nonce, DEFAULT_COOKIE_EXP)

	rt := ""
	rp := r.URL.Query().Get("redirect_uri")
	if rp != "" {
		rt = rp
	} else {
		ru := &url.URL{
			Scheme: "https",
			Host:   r.Host,
			Path:   path,
		}
		rt = ru.String()
	}

	req := url.Values{}
	req.Add("client_id", ctx.ClientId.String())
	req.Add("nonce", Sha256Hex(nonce))
	req.Add("redirect_uri", rt)
	req.Add("scope", "openid")
	req.Add("response_type", "id_token")

	u := URLMustParse(ctx.IdProviderAuthEndpoint.String())
	u.RawQuery = req.Encode()

	http.Redirect(w, r, u.String(), http.StatusFound)
}

func SetTokenCookieAndRedirect(w http.ResponseWriter, r *http.Request, token string) {
	SetSecureCookie(w, TOKEN_COOKIE_NAME, token, DEFAULT_COOKIE_EXP)

	q := r.URL.Query()
	q.Del("id_token")

	u := URLMustParse(r.URL.String())
	u.RawQuery = q.Encode()

	http.Redirect(w, r, u.String(), http.StatusFound)
}

func RequestHasForwardedUser(w http.ResponseWriter, r *http.Request) bool {
	_, ok := r.Header["X-Forwarded-User"]
	return ok
}

func RequestIsOverSecureChannel(w http.ResponseWriter, r *http.Request) bool {
	https, ok := r.Header["X-Forwarded-Proto"]
	if !ok || len(https) != 1 {
		glog.V(1).Infoln("Request does not contain X-Forwarded-Proto header")
		return false
	}

	if !CompareUpper(https[0], "HTTPS") {
		glog.V(1).Infoln("Request is not over HTTPS")
		return false
	}

	return true
}

func AuthProxyController(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context().Value("ProxyConfig").(*ProxyConfig)

	// Order matters in these checks!
	if RequestHasForwardedUser(w, r) {
		glog.Errorln("Request already has X-Forwarded-User")
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	if !RequestIsOverSecureChannel(w, r) {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

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

	rfpc, err := r.Cookie(RFP_COOKIE_NAME)
	if err != nil {
		glog.V(1).Infoln("No rfp cookie")
		if ctx.IsOptional {
			ctx.reverseProxy.ServeHTTP(w, r)
			return
		} else {
			RedirectToIdP(w, r, r.URL.Path)
			return
		}
	}

	if token := r.URL.Query().Get("id_token"); token != "" {
		if _, err := ctx.jwsValidator.Validate(token, rfpc.Value); err == nil {
			SetTokenCookieAndRedirect(w, r, token)
			return
		} else {
			glog.V(1).Infof("Querystring id_token invalid: %s", err)
		}
	}

	tokenc, err := r.Cookie(TOKEN_COOKIE_NAME)
	if err != nil {
		glog.V(1).Infoln("No token cookie")
		if ctx.IsOptional {
			ctx.reverseProxy.ServeHTTP(w, r)
			return
		} else {
			RedirectToIdP(w, r, r.URL.Path)
			return
		}
	}

	claims, err := ctx.jwsValidator.Validate(tokenc.Value, rfpc.Value)
	if err != nil {
		glog.Errorln("Token is invalid", err)
		if ctx.IsOptional {
			ctx.reverseProxy.ServeHTTP(w, r)
			return
		} else {
			RedirectToIdP(w, r, r.URL.Path)
			return
		}
	}

	r.Header["X-Forwarded-User"] = []string{claims.Subject}
	r.Header["X-OIDC-Token-Age"] = StringListFromInt(claims.Age())
	r.Header["X-OIDC-Token-Expires"] = StringListFromInt(int64(*claims.Expiry))
	r.Header["X-OIDC-UserInfo-Endpoint"] = []string{""} // TODO: Include this

	ctx.reverseProxy.ServeHTTP(w, r)
}

func LogoutController(w http.ResponseWriter, r *http.Request) {
	ExpireCookie(w, RFP_COOKIE_NAME)
	ExpireCookie(w, TOKEN_COOKIE_NAME)
	http.Redirect(w, r, "/", http.StatusFound)
}

// 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, error) {
	c := &ProxyConfig{}

	idpu := flag.String("idp", "", "URL for ID provider")
	mfam := flag.String("allow-mfa-methods", "", "Comma seperated list of allowed mfa methods")
	rmfa := flag.String("require-mfa-methods", "", "Comma seperated list of required mfa methods")
	cids := flag.String("client-id", "", "Client ID for proxy with IdP")

	flag.BoolVar(&c.IsOptional, "optional", false, "Allow proxying of unauthenticated calls")
	flag.BoolVar(&c.IsBootstrap, "bootstrap", false, "Allow running a proxy for the IdP itself")
	flag.BoolVar(&c.RequestMFA, "mfa", false, "Request user MFA authentication from IdP")

	flag.DurationVar(&c.MaxLiftetime, "max-lifetime", DEFAULT_MAX_LIFETIME, "Maximum allowed time from token issuance")
	flag.DurationVar(&c.ClockSkew, "clock-skew", DEFAULT_CLOCK_SKEW, "Allowable IdP clock skew relative to proxy")

	flag.StringVar(&c.UpstreamURL, "upstream", "", "URL of upstream service for which to proxy")
	flag.StringVar(&c.ListenOn, "listen", ":9992", "Optional port and ip on which to listen")
	flag.StringVar(&c.TrustedCACert, "ca", "", "Path to trusted CA certificate")

	flag.Parse()

	c.AllowedMFAMethods = strings.Split(*mfam, ",")
	c.RequiredMFAMethods = strings.Split(*rmfa, ",")

	if c.IsBootstrap {
		c.IsOptional = true
	}

	if _, err := os.Stat(c.TrustedCACert); os.IsNotExist(err) {
		return nil, errors.Errorf("CA certificate does not exist")
	}

	if cids == nil {
		return nil, errors.Errorf("Client ID is required")
	}

	if client_id, err := url.Parse(*cids); err != nil || client_id.Host == "" {
		return nil, errors.Errorf("Invalid client ID")
	} else {
		c.ClientId = client_id
	}

	if c.UpstreamURL == "" {
		return nil, errors.Errorf("Upstream URL is required")
	}

	if idpu == nil {
		return nil, errors.Errorf("IDP url is required")
	}

	if u, err := url.Parse(*idpu); err != nil {
		return nil, errors.WithStack(err)
	} else {
		c.IdProviderURL = u

		if h := HostFromURL(u.String()); c.IsBootstrap && (h != "localhost" && h != "127.0.0.1") {
			return nil, errors.Errorf("IdP must be set to localhost for bootstrap")
		}
	}

	return c, nil
}

func main() {
	cfg, err := parseConfig()
	if err != nil {
		glog.Fatalln("ParseConfig", err)
		return
	}

	hidp, err := NewCautiousHTTPClient(cfg.IsBootstrap)
	if err != nil {
		glog.Fatalln("Error building http client", err)
		return
	}

	idpc, err := FetchIdPConfig(hidp, cfg.IdProviderURL)
	if err != nil {
		glog.Fatalln("FetchIdPConfig:", err)
		return
	}

	cfg.IdProviderAuthEndpoint = idpc.AuthorizationEndpoint.AsURL()

	h, err := NewCautiousHTTPClient(false)
	if err != nil {
		glog.Fatalln("Error building http client", err)
		return
	}

	kf := NewJWKSFetcher(h, idpc.JwksUri.AsURL(), idpc.Issuer, cfg.TrustedCACert)

	cfg.jwsValidator = NewJWSValidator(&JWSValidationContext{
		KeyFetcher:   kf,
		Issuer:       idpc.Issuer,
		ClientId:     cfg.ClientId,
		ClockSkew:    cfg.ClockSkew,
		MaxLiftetime: cfg.MaxLiftetime,
	})

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

	http.HandleFunc("/.oidc/login", func(w http.ResponseWriter, r *http.Request) {
		RedirectToIdP(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)))
	})

	go http.ListenAndServe(cfg.ListenOn,
		handlers.LoggingHandler(os.Stdout, http.DefaultServeMux))

	// This has to happen last in-case we're boostrapping a proxy for the IdP itself
	if err := kf.Fetch(); err != nil {
		glog.Fatalln("FetchJWKS:", err)
		return
	} else {
		kf.Run()
	}
}