package main

import (
	"crypto"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/fcgi"
	"net/url"
	"os"

	"github.com/crewjam/saml"
	"github.com/crewjam/saml/logger"
)

// adapted from https://github.com/crewjam/saml/blob/main/example/idp/idp.go
func parsePEM(filename string) *pem.Block {
	raw, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}
	block, _ := pem.Decode(raw)
	if block == nil {
		panic("Could not decode PEM from " + filename)
	}
	return block
}

type Config struct {
	// The filesystem path to the private key to use for the IDP.
	// The key should be in PKCS8 PEM format.
	IDPKeyPath string `json:"idp_key_path"`
	// The filesystem path to the certificate to use for the IDP.
	// The certificate should be in PEM format.
	IDPCertPath string `json:"idp_cert_path"`
	// The name of the FCGI variable which has the user's username.
	UsernameKey string `json:"username_key"`
	// The name of the FCGI variable which has the user's email.
	EmailKey string `json:"email_key"`
	// The name of the FCGI variable which has the user's first name.
	GivenNameKey string `json:"given_name_key"`
	// The name of the FCGI variable which has the user's last name.
	SurnameKey string `json:"surname_key"`
	// The name of the FCGI variable which has the user's groups.
	// The groups should be concatenated together into one string
	// separated by semicolons, e.g. "group1;group2".
	GroupsKey string `json:"groups_key"`
	// The base URL for the IDP. The following endpoints will be
	// derived from it:
	// Metadata: baseURL + "/metadata"
	// SSO login: baseURL + "/sso"
	BaseURL string `json:"base_url"`
	// An array of filepaths or HTTP URLs of the service providers'
	// metadata XML files.
	SPMetadataPaths []string `json:"sp_metadata_paths"`
	// The filesystem path of the Unix domain socket on which this
	// program will listen for incoming FCGI connections.
	SocketPath string `json:"socket_path"`
}

func newIDP(cfg *Config) *saml.IdentityProvider {
	key := func() crypto.PrivateKey {
		block := parsePEM(cfg.IDPKeyPath)
		k, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			panic(err)
		}
		return k
	}()
	cert := func() *x509.Certificate {
		block := parsePEM(cfg.IDPCertPath)
		c, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			panic(err)
		}
		return c
	}()
	sessionProvider := &BasicSessionProvider{
		UsernameKey:  cfg.UsernameKey,
		EmailKey:     cfg.EmailKey,
		GivenNameKey: cfg.GivenNameKey,
		SurnameKey:   cfg.SurnameKey,
		GroupsKey:    cfg.GroupsKey,
	}
	serviceProviderProvider := NewServiceProviderProvider(cfg.SPMetadataPaths)
	metadataURL, err := url.Parse(cfg.BaseURL + "/metadata")
	if err != nil {
		panic(err)
	}
	ssoURL, err := url.Parse(cfg.BaseURL + "/sso")
	if err != nil {
		panic(err)
	}
	return &saml.IdentityProvider{
		Key:                     key,
		Certificate:             cert,
		Logger:                  logger.DefaultLogger,
		MetadataURL:             *metadataURL,
		SSOURL:                  *ssoURL,
		SessionProvider:         sessionProvider,
		ServiceProviderProvider: serviceProviderProvider,
	}
}

func newHTTPHandler(idp *saml.IdentityProvider) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc(idp.MetadataURL.Path, idp.ServeMetadata)
	mux.HandleFunc(idp.SSOURL.Path, idp.ServeSSO)
	return mux
}

func newListener(cfg *Config) net.Listener {
	var err error
	sockPath := cfg.SocketPath
	if _, err = os.Stat(sockPath); err == nil {
		err = os.Remove(sockPath)
		if err != nil {
			panic(err)
		}
	}
	ln, err := net.Listen("unix", sockPath)
	if err != nil {
		panic(err)
	}
	os.Chmod(sockPath, 0777)
	return ln
}

func main() {
	// Allow the key path to be overriden via a CLI flag so that we
	// can use systemd's LoadCredential feature
	idpKeyPathPtr := flag.String("k", "", "IDP key path")
	configPathPtr := flag.String("c", "./config.json", "config file path")
	flag.Parse()

	cfgBytes, err := ioutil.ReadFile(*configPathPtr)
	if err != nil {
		panic(err)
	}
	cfg := Config{}
	err = json.Unmarshal(cfgBytes, &cfg)
	if err != nil {
		panic(err)
	}
	if *idpKeyPathPtr != "" {
		cfg.IDPKeyPath = *idpKeyPathPtr
	}

	idp := newIDP(&cfg)
	handler := newHTTPHandler(idp)
	listener := newListener(&cfg)

	err = fcgi.Serve(listener, handler)
	if err != nil {
		panic(err)
	}
}