package pki

// Copyright 2018 Sergiusz Bazanski <q3k@hackerspace.pl>
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
// ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR
// IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"strings"

	"github.com/golang/glog"
	"golang.org/x/net/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"
)

var (
	flagCAPath          string
	flagCertificatePath string
	flagKeyPath         string
	flagPKICluster      string
	flagPKIRealm        string
	flagPKIDisable      bool

	// Enable logging HSPKI info into traces
	Trace = true
	// Enable logging HSPKI info into glog
	Log = false
)

const (
	ctxKeyClientInfo = "hspki-client-info"
)

func init() {
	flag.StringVar(&flagCAPath, "hspki_tls_ca_path", "pki/ca.pem", "Path to PKI CA certificate")
	flag.StringVar(&flagCertificatePath, "hspki_tls_certificate_path", "pki/service.pem", "Path to PKI service certificate")
	flag.StringVar(&flagKeyPath, "hspki_tls_key_path", "pki/service-key.pem", "Path to PKI service private key")
	flag.StringVar(&flagPKICluster, "hspki_cluster", "local.hswaw.net", "FQDN of cluster on which this service runs")
	flag.StringVar(&flagPKIRealm, "hspki_realm", "hswaw.net", "Cluster realm (top level from which we accept foreign cluster certs)")
	flag.BoolVar(&flagPKIDisable, "hspki_disable", false, "Disable PKI entirely (insecure!)")
}

func maybeTrace(ctx context.Context, f string, args ...interface{}) {
	if Log {
		glog.Infof(f, args...)
	}

	if !Trace {
		return
	}

	tr, ok := trace.FromContext(ctx)
	if !ok {
		if !Log {
			fmtd := fmt.Sprintf(f, args...)
			glog.Info("[no trace] %v", fmtd)
		}
		return
	}
	tr.LazyPrintf(f, args...)
}

func parseClientName(name string) (*ClientInfo, error) {
	if !strings.HasSuffix(name, "."+flagPKIRealm) {
		return nil, fmt.Errorf("invalid realm")
	}

	inRealm := strings.TrimSuffix(name, "."+flagPKIRealm)

	special := []string{"person", "external"}

	for _, s := range special {
		// Special case for people running jobs from workstations, or for non-cluster services.
		if strings.HasSuffix(inRealm, "."+s) {
			asPerson := strings.TrimSuffix(inRealm, "."+s)
			parts := strings.Split(asPerson, ".")
			if len(parts) != 1 {
				return nil, fmt.Errorf("invalid person fqdn")
			}
			return &ClientInfo{
				Cluster:   fmt.Sprintf("%s.%s", s, flagPKIRealm),
				Principal: parts[0],
				Job:       "",
			}, nil
		}
	}

	parts := strings.Split(inRealm, ".")
	if len(parts) != 4 {
		return nil, fmt.Errorf("invalid job/principal format for in-cluster")
	}
	if parts[2] != "svc" {
		return nil, fmt.Errorf("can only refer to services within cluster")
	}
	clusterShort := parts[3]

	return &ClientInfo{
		Cluster:   fmt.Sprintf("%s.%s", clusterShort, flagPKIRealm),
		Principal: fmt.Sprintf("%s.svc", parts[1]),
		Job:       parts[0],
	}, nil
}

func withPKIInfo(ctx context.Context, c *ClientInfo) context.Context {
	maybeTrace(ctx, "HSPKI: Applying ClientInfo: %s", c.String())
	return context.WithValue(ctx, ctxKeyClientInfo, c)
}

func grpcInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
	peer, ok := peer.FromContext(ctx)
	if !ok {
		maybeTrace(ctx, "HSPKI: Could not establish identity of peer.")
		return nil, status.Errorf(codes.PermissionDenied, "no peer info")
	}

	authInfo, ok := peer.AuthInfo.(credentials.TLSInfo)
	if !ok {
		maybeTrace(ctx, "HSPKI: Could not establish TLS identity of peer.")
		return nil, status.Errorf(codes.PermissionDenied, "no TLS certificate presented")
	}

	chains := authInfo.State.VerifiedChains
	if len(chains) != 1 {
		maybeTrace(ctx, "HSPKI: No trusted chains found.")
		return nil, status.Errorf(codes.PermissionDenied, "no trusted TLS certificate presented")
	}

	chain := chains[0]

	certDNs := make([]string, len(chain))
	for i, cert := range chain {
		certDNs[i] = cert.Subject.String()
	}
	maybeTrace(ctx, "HSPKI: Trust chain: %s", strings.Join(certDNs, ", "))

	clientInfo, err := parseClientName(chain[0].Subject.CommonName)
	if err != nil {
		maybeTrace(ctx, "HSPKI: Invalid CN %q: %v", chain[0].Subject.CommonName, err)
		return nil, status.Errorf(codes.PermissionDenied, "invalid TLS CN format")
	}
	ctx = withPKIInfo(ctx, clientInfo)
	return handler(ctx, req)
}

// ClientInfo contains information about the HSPKI authentication data of the
// gRPC client that has made the request.
type ClientInfo struct {
	Cluster   string
	Principal string
	Job       string
}

// String returns a human-readable representation of the ClientInfo in the
// form "job=foo, principal=bar.svc, cluster=baz.hswaw.net".
func (c *ClientInfo) String() string {
	return fmt.Sprintf("job=%q, principal=%q, cluster=%q", c.Job, c.Principal, c.Cluster)
}

// Person returns a reference to a person's ID if the ClientInfo describes a person.
// Otherwise, it returns an empty string.
func (c *ClientInfo) Person() string {
	if c.Cluster != fmt.Sprintf("person.%s", flagPKIRealm) {
		return ""
	}
	return c.Principal
}

// ClientInfoFromContext returns ClientInfo from a gRPC service context.
func ClientInfoFromContext(ctx context.Context) *ClientInfo {
	v := ctx.Value(ctxKeyClientInfo)
	if v == nil {
		return nil
	}
	ci, ok := v.(*ClientInfo)
	if !ok {
		return nil
	}
	return ci
}

// WithServerHSPKI is a grpc.ServerOptions array that ensures that the gRPC server:
// - runs with HSPKI TLS Service Certificate
// - rejects all non_HSPKI compatible requests
// - injects ClientInfo into the service context, which can be later retrieved
//   using ClientInfoFromContext
func WithServerHSPKI() []grpc.ServerOption {
	if !flag.Parsed() {
		glog.Exitf("WithServerHSPKI called before flag.Parse!")
	}
	if flagPKIDisable {
		return []grpc.ServerOption{}
	}

	loc, err := loadCredentials()
	if err != nil {
		glog.Exitf("WithServerHSPKI: loadCredentials: %v", err)
	}

	serverCert, err := tls.X509KeyPair(loc.cert, loc.key)
	if err != nil {
		glog.Exitf("WithServerHSPKI: cannot load service certificate/key: %v", err)
	}

	certPool := x509.NewCertPool()
	if ok := certPool.AppendCertsFromPEM(loc.ca); !ok {
		glog.Exitf("WithServerHSPKI: cannot use CA certificate")
	}

	creds := grpc.Creds(credentials.NewTLS(&tls.Config{
		ClientAuth:   tls.RequireAndVerifyClientCert,
		Certificates: []tls.Certificate{serverCert},
		ClientCAs:    certPool,
	}))

	interceptor := grpc.UnaryInterceptor(grpcInterceptor)

	return []grpc.ServerOption{creds, interceptor}
}

type ClientHSPKIOption func(c *tls.Config)

func OverrideServerName(name string) ClientHSPKIOption {
	return func(c *tls.Config) {
		c.ServerName = name
	}
}

func WithClientHSPKI(opts ...ClientHSPKIOption) grpc.DialOption {
	if !flag.Parsed() {
		glog.Exitf("WithServerHSPKI called before flag.Parse!")
	}
	if flagPKIDisable {
		return grpc.WithInsecure()
	}

	loc, err := loadCredentials()
	if err != nil {
		glog.Exitf("WithServerHSPKI: loadCredentials: %v", err)
	}

	certPool := x509.NewCertPool()
	if ok := certPool.AppendCertsFromPEM(loc.ca); !ok {
		glog.Exitf("WithServerHSPKI: cannot use CA certificate")
	}

	clientCert, err := tls.X509KeyPair(loc.cert, loc.key)
	if err != nil {
		glog.Exitf("WithClientHSPKI: cannot load service certificate/key: %v", err)
	}

	config := &tls.Config{
		Certificates: []tls.Certificate{clientCert},
		RootCAs:      certPool,
	}

	for _, opt := range opts {
		opt(config)
	}

	creds := credentials.NewTLS(config)
	return grpc.WithTransportCredentials(creds)
}