// Copyright 2021 The Witness Contributors
//
// 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 gcpiit

import (
	"crypto"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/in-toto/go-witness/attestation"
	"github.com/in-toto/go-witness/attestation/jwt"
	"github.com/in-toto/go-witness/cryptoutil"
	"github.com/in-toto/go-witness/log"
	"github.com/invopop/jsonschema"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	Name    = "gcp-iit"
	Type    = "https://witness.dev/attestations/gcp-iit/v0.1"
	RunType = attestation.PreMaterialRunType

	jwksUrl                      = "https://www.googleapis.com/oauth2/v3/certs"
	defaultIdentityTokenHost     = "metadata.google.internal"
	identityTokenURLPathTemplate = "/computeMetadata/v1/instance/service-accounts/%s/identity"
	identityTokenAudience        = "witness-node-attestor" //nolint: gosec // false positive
	defaultServiceAccount        = "default"
	TokenUrl                     = "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?audience=witness-node-attestor&format=full&licenses=TRUE"
	InstanceMetadataUrl          = "http://metadata.google.internal/computeMetadata/v1/instance/"
	InstanceAttributesUrl        = "http://metadata.google.internal/computeMetadata/v1/instance/attributes/"
	ProjectMetadataUrl           = "http://metadata.google.internal/computeMetadata/v1/project/"
)

// This is a hacky way to create a compile time error in case the attestor
// doesn't implement the expected interfaces.
var (
	_ attestation.Attestor  = &Attestor{}
	_ attestation.Subjecter = &Attestor{}
)

func init() {
	attestation.RegisterAttestation(Name, Type, RunType, func() attestation.Attestor {
		return New()
	})
}

type ErrNotGCPIIT struct{}

func (e ErrNotGCPIIT) Error() string {
	return "not a GCP IIT JWT"
}

type Attestor struct {
	JWT                       *jwt.Attestor `json:"jwt"`
	ProjectID                 string        `json:"project_id"`
	ProjectNumber             string        `json:"project_number"`
	InstanceZone              string        `json:"zone"`
	InstanceID                string        `json:"instance_id"`
	InstanceHostname          string        `json:"instance_hostname"`
	InstanceCreationTimestamp string        `json:"instance_creation_timestamp"`
	InstanceConfidentiality   string        `json:"instance_confidentiality"`
	LicenceID                 []string      `json:"licence_id"`
	ClusterName               string        `json:"cluster_name"`
	ClusterUID                string        `json:"cluster_uid"`
	ClusterLocation           string        `json:"cluster_location"`

	isWorkloadIdentity bool
}

func New() *Attestor {
	return &Attestor{}
}

func (a *Attestor) Name() string {
	return Name
}

func (a *Attestor) Type() string {
	return Type
}

func (a *Attestor) RunType() attestation.RunType {
	return RunType
}

func (a *Attestor) Schema() *jsonschema.Schema {
	// NOTE: This isn't ideal. For some reason the reflect function is return an empty schema when passing in `p`
	// TODO: Fix this later
	schema := jsonschema.Reflect(&a)
	schema.Definitions["Attestor"].Properties.Set("jwt", jsonschema.Reflect(&a.JWT))
	return schema
}

func (a *Attestor) Attest(ctx *attestation.AttestationContext) error {
	tokenURL := identityTokenURL(defaultIdentityTokenHost, defaultServiceAccount)
	identityToken, err := getMetadata(tokenURL)
	if err != nil {
		// status.Errorf does not support %w directive
		return status.Errorf(codes.Internal, "unable to retrieve valid identity token: %v", err)
	}

	it := string(identityToken)

	a.JWT = jwt.New(jwt.WithToken(it), jwt.WithJWKSUrl(jwksUrl))
	if err := a.JWT.Attest(ctx); err != nil {
		return err
	}

	if a.JWT.Claims["google"] == nil {
		a.isWorkloadIdentity = true
	}

	if !a.isWorkloadIdentity {
		googClaim := a.JWT.Claims["google"].(map[string]interface{})
		a.ProjectID = googClaim["project_id"].(string)
		a.ProjectNumber = googClaim["project_number"].(string)
		a.InstanceZone = googClaim["zone"].(string)
		a.InstanceID = googClaim["instance_id"].(string)
		a.InstanceHostname = googClaim["instance_name"].(string)
		a.InstanceCreationTimestamp = googClaim["instance_creation_timestamp"].(string)
		a.InstanceConfidentiality = googClaim["instance_confidentiality"].(string)
		a.LicenceID = googClaim["licence_id"].([]string)
	} else {
		a.getInstanceData()
	}

	return nil
}

func (a *Attestor) getInstanceData() {
	endpoints := map[string]string{
		"hostname":         InstanceMetadataUrl + "hostname",
		"id":               InstanceMetadataUrl + "id",
		"zone":             InstanceMetadataUrl + "zone",
		"cluster-name":     InstanceMetadataUrl + "attributes/cluster-name",
		"cluster-uid":      InstanceMetadataUrl + "attributes/cluster-uid",
		"cluster-location": InstanceMetadataUrl + "attributes/cluster-location",
		"project-id":       InstanceMetadataUrl + "project/project-id",
		"project-number":   InstanceMetadataUrl + "project/numeric-project-id",
	}

	metadata := make(map[string]string)

	for k, v := range endpoints {
		data, err := getMetadata(v)
		if err != nil {
			log.Warnf("failed to retrieve gcp metadata from %v: %w", v, err)
			continue
		}
		metadata[k] = string(data)
	}

	a.ClusterName = metadata["cluster-name"]
	a.ClusterUID = metadata["cluster-uid"]
	a.ClusterLocation = metadata["cluster-location"]
	a.InstanceHostname = metadata["hostname"]
	a.InstanceID = metadata["id"]
	a.InstanceZone = metadata["zone"]

	projID, projNum, err := parseJWTProjectInfo(a.JWT)
	if err != nil {
		log.Warnf("unable to parse gcp project info from JWT: %w\n", err)
	}

	a.ProjectID = projID
	a.ProjectNumber = projNum
}

func (a *Attestor) Subjects() map[string]cryptoutil.DigestSet {
	subjects := make(map[string]cryptoutil.DigestSet)
	hashes := []cryptoutil.DigestValue{{Hash: crypto.SHA256}}
	if ds, err := cryptoutil.CalculateDigestSetFromBytes([]byte(a.InstanceID), hashes); err == nil {
		subjects[fmt.Sprintf("instanceid:%v", a.InstanceID)] = ds
	} else {
		log.Debugf("(attestation/gcp) failed to record gcp instanceid subject: %w", err)
	}

	if ds, err := cryptoutil.CalculateDigestSetFromBytes([]byte(a.InstanceHostname), hashes); err == nil {
		subjects[fmt.Sprintf("instancename:%v", a.InstanceHostname)] = ds
	} else {
		log.Debugf("(attestation/gcp) failed to record gcp instancename subject: %w", err)
	}

	if ds, err := cryptoutil.CalculateDigestSetFromBytes([]byte(a.ProjectID), hashes); err == nil {
		subjects[fmt.Sprintf("projectid:%v", a.ProjectID)] = ds
	} else {
		log.Debugf("(attestation/gcp) failed to record gcp projectid subject: %w", err)
	}

	if ds, err := cryptoutil.CalculateDigestSetFromBytes([]byte(a.ProjectNumber), hashes); err == nil {
		subjects[fmt.Sprintf("projectnumber:%v", a.ProjectNumber)] = ds
	} else {
		log.Debugf("(attestation/gcp) failed to record gcp projectnumber subject: %w", err)
	}

	if ds, err := cryptoutil.CalculateDigestSetFromBytes([]byte(a.ClusterUID), hashes); err == nil {
		subjects[fmt.Sprintf("clusteruid:%v", a.ClusterUID)] = ds
	} else {
		log.Debugf("(attestation/gcp) failed to record gcp clusteruid subject: %w", err)
	}

	return subjects
}

func getMetadata(url string) ([]byte, error) {
	client := &http.Client{}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("unexpected status code for route: %d", resp.StatusCode)
	}

	bytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return bytes, nil
}

// identityTokenURL creates the URL to find an instance identity document given the
// host of the GCP metadata server and the service account the instance is running as.
func identityTokenURL(host, serviceAccount string) string {
	query := url.Values{}
	query.Set("audience", identityTokenAudience)
	query.Set("format", "full")
	url := &url.URL{
		Scheme:   "http",
		Host:     host,
		Path:     fmt.Sprintf(identityTokenURLPathTemplate, serviceAccount),
		RawQuery: query.Encode(),
	}
	return url.String()
}

func parseJWTProjectInfo(jwt *jwt.Attestor) (string, string, error) {
	if jwt.Claims["email"] == nil {
		return "", "", fmt.Errorf("unable to find email claim")
	}

	email := jwt.Claims["email"].(string)

	stings := strings.Split(email, "@")
	if len(stings) != 2 {
		return "", "", fmt.Errorf("unable to parse email: %s", email)
	}

	domain := stings[1]

	projectInfo := strings.Split(domain, ".")[0]
	projectInfoSplit := strings.Split(projectInfo, "-")
	projectID := projectInfoSplit[len(projectInfoSplit)-1]
	projectName := strings.Join(projectInfoSplit[:len(projectInfoSplit)-1], "-")

	return projectID, projectName, nil
}
