/*
Copyright 2021 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package podsecurity

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"

	// install conversions for types we need to convert
	_ "k8s.io/kubernetes/pkg/apis/apps/install"
	_ "k8s.io/kubernetes/pkg/apis/batch/install"
	_ "k8s.io/kubernetes/pkg/apis/core/install"

	appsv1 "k8s.io/api/apps/v1"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apiserver/pkg/admission"
	genericadmissioninit "k8s.io/apiserver/pkg/admission/initializer"
	"k8s.io/apiserver/pkg/audit"
	"k8s.io/apiserver/pkg/warning"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"k8s.io/component-base/featuregate"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/apis/apps"
	"k8s.io/kubernetes/pkg/apis/batch"
	"k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/features"
	podsecurityadmission "k8s.io/pod-security-admission/admission"
	podsecurityconfigloader "k8s.io/pod-security-admission/admission/api/load"
	podsecurityadmissionapi "k8s.io/pod-security-admission/api"
	"k8s.io/pod-security-admission/policy"
)

// PluginName is a string with the name of the plugin
const PluginName = "PodSecurity"

// Register registers a plugin
func Register(plugins *admission.Plugins) {
	plugins.Register(PluginName, func(reader io.Reader) (admission.Interface, error) {
		return newPlugin(reader)
	})
}

// Plugin holds state for and implements the admission plugin.
type Plugin struct {
	*admission.Handler

	enabled               bool
	inspectedFeatureGates bool

	client          kubernetes.Interface
	namespaceLister corev1listers.NamespaceLister
	podLister       corev1listers.PodLister

	delegate *podsecurityadmission.Admission
}

var _ admission.ValidationInterface = &Plugin{}
var _ genericadmissioninit.WantsExternalKubeInformerFactory = &Plugin{}
var _ genericadmissioninit.WantsExternalKubeClientSet = &Plugin{}

// newPlugin creates a new admission plugin.
func newPlugin(reader io.Reader) (*Plugin, error) {
	config, err := podsecurityconfigloader.LoadFromReader(reader)
	if err != nil {
		return nil, err
	}

	evaluator, err := policy.NewEvaluator(policy.DefaultChecks())
	if err != nil {
		return nil, fmt.Errorf("could not create PodSecurityRegistry: %w", err)
	}

	return &Plugin{
		Handler: admission.NewHandler(admission.Create, admission.Update),
		delegate: &podsecurityadmission.Admission{
			Configuration:    config,
			Evaluator:        evaluator,
			Metrics:          nil, // TODO: wire to default prometheus metrics
			PodSpecExtractor: podsecurityadmission.DefaultPodSpecExtractor{},
		},
	}, nil
}

// SetExternalKubeInformerFactory registers an informer
func (p *Plugin) SetExternalKubeInformerFactory(f informers.SharedInformerFactory) {
	namespaceInformer := f.Core().V1().Namespaces()
	p.namespaceLister = namespaceInformer.Lister()
	p.podLister = f.Core().V1().Pods().Lister()
	p.SetReadyFunc(namespaceInformer.Informer().HasSynced)
	p.updateDelegate()
}

// SetExternalKubeClientSet sets the plugin's client
func (p *Plugin) SetExternalKubeClientSet(client kubernetes.Interface) {
	p.client = client
	p.updateDelegate()
}

func (p *Plugin) updateDelegate() {
	// return early if we don't have what we need to set up the admission delegate
	if p.namespaceLister == nil {
		return
	}
	if p.podLister == nil {
		return
	}
	if p.client == nil {
		return
	}
	p.delegate.PodLister = podsecurityadmission.PodListerFromInformer(p.podLister)
	p.delegate.NamespaceGetter = podsecurityadmission.NamespaceGetterFromListerAndClient(p.namespaceLister, p.client)
}

func (c *Plugin) InspectFeatureGates(featureGates featuregate.FeatureGate) {
	c.enabled = featureGates.Enabled(features.PodSecurity)
	c.inspectedFeatureGates = true
}

// ValidateInitialization ensures all required options are set
func (p *Plugin) ValidateInitialization() error {
	if !p.inspectedFeatureGates {
		return fmt.Errorf("%s did not see feature gates", PluginName)
	}
	if err := p.delegate.CompleteConfiguration(); err != nil {
		return fmt.Errorf("%s configuration error: %w", PluginName, err)
	}
	if err := p.delegate.ValidateConfiguration(); err != nil {
		return fmt.Errorf("%s invalid: %w", PluginName, err)
	}
	return nil
}

var (
	applicableResources = map[schema.GroupResource]bool{
		corev1.Resource("pods"):       true,
		corev1.Resource("namespaces"): true,
	}
)

func (p *Plugin) Validate(ctx context.Context, a admission.Attributes, o admission.ObjectInterfaces) error {
	if !p.enabled {
		return nil
	}
	gr := a.GetResource().GroupResource()
	if !applicableResources[gr] && !p.delegate.PodSpecExtractor.HasPodSpec(gr) {
		return nil
	}

	a = &lazyConvertingAttributes{Attributes: a}

	result := p.delegate.Validate(ctx, a)
	for _, w := range result.Warnings {
		warning.AddWarning(ctx, "", w)
	}
	for k, v := range result.AuditAnnotations {
		audit.AddAuditAnnotation(ctx, podsecurityadmissionapi.AuditAnnotationPrefix+k, v)
	}
	if !result.Allowed {
		if result.Result != nil && len(result.Result.Message) > 0 {
			// TODO: use code/reason/etc from status
			return admission.NewForbidden(a, errors.New(result.Result.Message))
		}
		return admission.NewForbidden(a, errors.New("Not allowed by PodSecurity"))
	}
	return nil
}

type lazyConvertingAttributes struct {
	admission.Attributes

	convertObjectOnce sync.Once
	convertedObject   runtime.Object

	convertOldObjectOnce sync.Once
	convertedOldObject   runtime.Object
}

func (l *lazyConvertingAttributes) GetObject() runtime.Object {
	l.convertObjectOnce.Do(func() {
		obj, err := convert(l.Attributes.GetObject())
		if err != nil {
			utilruntime.HandleError(err)
		}
		l.convertedObject = obj
	})
	return l.convertedObject
}
func (l *lazyConvertingAttributes) GetOldObject() runtime.Object {
	l.convertOldObjectOnce.Do(func() {
		obj, err := convert(l.Attributes.GetOldObject())
		if err != nil {
			utilruntime.HandleError(err)
		}
		l.convertedOldObject = obj
	})
	return l.convertedOldObject
}

func convert(in runtime.Object) (runtime.Object, error) {
	var out runtime.Object
	switch in.(type) {
	case *core.Namespace:
		out = &corev1.Namespace{}
	case *core.Pod:
		out = &corev1.Pod{}
	case *core.ReplicationController:
		out = &corev1.ReplicationController{}
	case *core.PodTemplate:
		out = &corev1.PodTemplate{}
	case *apps.ReplicaSet:
		out = &appsv1.ReplicaSet{}
	case *apps.Deployment:
		out = &appsv1.Deployment{}
	case *apps.StatefulSet:
		out = &appsv1.StatefulSet{}
	case *apps.DaemonSet:
		out = &appsv1.DaemonSet{}
	case *batch.Job:
		out = &batchv1.Job{}
	case *batch.CronJob:
		out = &batchv1.CronJob{}
	default:
		return in, fmt.Errorf("unexpected type %T", in)
	}
	if err := legacyscheme.Scheme.Convert(in, out, nil); err != nil {
		return in, err
	}
	return out, nil
}