// Copyright 2023 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package bcid

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"os/exec"
	"runtime"

	spb "github.com/in-toto/attestation/go/v1"
	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/cipd/common"
	"go.chromium.org/luci/common/logging"
	"golang.org/x/oauth2"
	"google.golang.org/protobuf/encoding/protojson"
	structpb "google.golang.org/protobuf/types/known/structpb"
)

const (
	ProdKeyID     = "projects/fuchsia-bcid/locations/global/keyRings/bcid/cryptoKeys/attestation_signer/cryptoKeyVersions/1"
	TestKeyID     = "projects/fuchsia-bcid-dev/locations/global/keyRings/tests/cryptoKeys/recipe_wrapper_attestation_signer/cryptoKeyVersions/1"
	StmtType      = "https://in-toto.io/Statement/v1"
	PredicateType = "https://slsa.dev/provenance/v1"
	sciloEndpoint = "https://bcidsoftwareverifier-pa.googleapis.com/v1/software-artifact-verification-requests"
)

type attestOutputter struct{}

func (*attestOutputter) CombinedOutput(ctx context.Context, stdin []byte, args ...string) ([]byte, error) {
	cmd := exec.CommandContext(ctx, "attestation_tool", args...)
	if stdin != nil {
		pipe, err := cmd.StdinPipe()
		if err != nil {
			return nil, fmt.Errorf("unable to pipe to stdin: %v", err)
		}
		pipe.Write(stdin)
		// This close must occur before CombinedOutput() is called otherwise the command will wait forever.
		pipe.Close()
	}

	logging.Infof(ctx, "Executing %+v\n", cmd)
	return cmd.CombinedOutput()
}

func CanAttestPlatform() bool {
	return runtime.GOOS == "linux" && runtime.GOARCH == "amd64"
}

func Attest(ctx context.Context, stmt *spb.Statement, keyID string) (*InTotoBundle, error) {
	if !CanAttestPlatform() {
		return nil, fmt.Errorf("platform %v/%v is unsupported for attestation", runtime.GOOS, runtime.GOARCH)
	}
	return attestWithOutputter(ctx, &attestOutputter{}, stmt, keyID)
}

type CombinedOutputer interface {
	CombinedOutput(context.Context, []byte, ...string) ([]byte, error)
}

func attestWithOutputter(ctx context.Context, cmd CombinedOutputer, stmt *spb.Statement, keyID string) (*InTotoBundle, error) {
	stmtJSON, err := protojson.Marshal(stmt)
	if err != nil {
		return nil, err
	}

	out, err := cmd.CombinedOutput(ctx, stmtJSON, "sign_intoto_statement",
		"--input", "-",
		"--key_id", keyID,
		"--output", "-")
	if err != nil {
		return nil, fmt.Errorf("attestation_tool failed: %v / %q", err, out)
	}
	return UnmarshalBundle(bytes.NewReader(out))
}

type Signature struct {
	KeyID string `json:"keyid"`
	Sig   string `json:"sig"`
}

// InToToBundle represents the attestation bundle.
// TODO(https://github.com/in-toto/attestation/issues/280): This struct should ideally be provided upstream.
type InTotoBundle struct {
	ContentType string      `json:"payloadType"`
	Payload     string      `json:"payload"`
	Signatures  []Signature `json:"signatures"`
}

func UnmarshalBundle(data io.Reader) (*InTotoBundle, error) {
	all, err := io.ReadAll(data)
	if err != nil {
		return nil, fmt.Errorf("can't read input data: %v", err)
	}

	var itt InTotoBundle
	if err := json.Unmarshal(all, &itt); err != nil {
		return nil, fmt.Errorf("unable to pass in-toto payload: %v", err)
	}

	return &itt, nil
}

type Subject struct {
	Pkg  *common.Pin
	Data io.Reader
}

func NewStmt(subject *Subject) (*spb.Statement, error) {
	if subject == nil {
		return nil, errors.New("subject cannot be nil")
	}
	h := sha256.New()
	if _, err := io.Copy(h, subject.Data); err != nil {
		return nil, err
	}

	// The statement proto uses this struct proto which is dynamically typed.
	// It then uses reflection to try and figure out what it all is.
	// The only real way to work through this is by using raw types, not
	// real Go structs.
	// It's possible to just dump raw JSON here, but that feels even less
	// idomatic than this does.
	// Unless tagged as required, fields are optional.
	// Any changes to this specification should be reflected in the external
	// documentation at //slsa/buildConfig/v1
	// Do not remove fields once populated unless BCID For Software team is
	// consulted.
	pred, err := structpb.NewStruct(map[string]any{
		// Required.
		"buildDefinition": map[string]any{
			// Required.
			"buildType": "https://fuchsia.googlesource.com/infra/infra/slsa/buildConfig/v1",
			// Required, but fields within may be best-effort.
			// The spec is unclear whether this field can be included empty.
			"externalParameters": map[string]any{
				"buildConfigSource": map[string]any{
					"repository": "git+https://fuchsia.googlesource.com/infra/recipes",
					"path":       "recipes.py",
				},
			},
		},
		// Required.
		"runDetails": map[string]any{
			// Required.
			"builder": map[string]any{
				// Required.
				"id": "//bcid.corp.google.com/builders/luci/fuchsia/l1",
				// TODO(cflewis): Add the version here.
			},
		},
	})
	if err != nil {
		return nil, err
	}

	return &spb.Statement{
		Type: StmtType,
		Subject: []*spb.ResourceDescriptor{
			{
				Name: subject.Pkg.PackageName,
				// cipd_instance_id is required by SCILo to understand what is being uploaded.
				// This key isn't arbitrary chosen and _mustn't be changed_.
				Digest: map[string]string{"cipd_instance_id": subject.Pkg.InstanceID},
			},
		},
		PredicateType: PredicateType,
		Predicate:     pred,
	}, nil
}

type sciloContext struct {
	VerificationPurpose  string `json:"verificationPurpose"`
	EnforcementPointName string `json:"enforcementPointName"`
	OccurrenceStage      string `json:"occurrenceStage"`
}

type sciloArtifactInfo struct {
	Digests map[string]string `json:"digests"`
	// Attestations are _strings_ of InTotoBundles.
	// Passing bundles as JSON is not accepted by the endpoint.
	Attestations []string `json:"attestations"`
}

type sciloReq struct {
	Context          *sciloContext      `json:"context"`
	ResourceToVerify string             `json:"resourceToVerify"`
	ArtifactInfo     *sciloArtifactInfo `json:"artifactInfo"`
}

type SCILoResp struct {
	Allowed             bool   `json:"allowed"`
	RejectionMessage    string `json:"rejectionMessage"`
	VerificationSummary string `json:"verificationSummary"`
}

// UploadBundle uploads the bundle to BCID using default values.
func UploadBundle(ctx context.Context, stmt *spb.Statement, bundle *InTotoBundle) (*SCILoResp, error) {
	ts, err := LUCITokenSource(ctx)
	if err != nil {
		return nil, err
	}
	return UploadBundleToEndpoint(ctx, sciloEndpoint, ts, stmt, bundle)
}

func UploadBundleToEndpoint(ctx context.Context,
	endpoint string, tokenSource oauth2.TokenSource, stmt *spb.Statement, bundle *InTotoBundle) (*SCILoResp, error) {
	token, err := tokenSource.Token()
	if err != nil {
		return nil, fmt.Errorf("couldn't get auth token: %v", err)
	}
	if subjs := len(stmt.GetSubject()); subjs != 1 {
		return nil, fmt.Errorf("expected 1 subject in statement, got %v", subjs)
	}
	subj := stmt.GetSubject()[0]

	jsonBundle, err := json.Marshal(bundle)
	if err != nil {
		return nil, fmt.Errorf("couldn't marshal bundle to JSON: %v", err)
	}

	sReq := &sciloReq{
		// These values are documented at go/scilo-server#artifact-occurrence-stage
		Context: &sciloContext{
			VerificationPurpose: "VERIFY_FOR_ENFORCEMENT",
			// This is an arbitrary name to identify what is attesting.
			EnforcementPointName: "fuchsia-recipe-wrapper",
			// This means that the attestation has been made after the build has completed
			// and will no longer be modified.
			OccurrenceStage: "AS_VERIFIED",
		},
		ResourceToVerify: "cipd_package://" + subj.GetName(),
		ArtifactInfo: &sciloArtifactInfo{
			Digests:      subj.GetDigest(),
			Attestations: []string{string(jsonBundle)},
		},
	}
	jsonReq, err := json.Marshal(sReq)
	if err != nil {
		return nil, fmt.Errorf("couldn't marshal SCILo Request to JSON: %v", err)
	}
	logging.Infof(ctx, "Uploading attestations to SCILo:\n%s", jsonBundle)

	httpReq, _ := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(jsonReq))
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "application/json")
	httpReq.Header.Set("Authorization", fmt.Sprintf("Bearer %v", token.AccessToken))
	// The dump needs to happen _before_ the httpclient.Do() call.
	// Do() probably consumes the request and then dump is empty.
	// DumpRequestOut() will copy back the data it consumes so the
	// Do() call works correctly.
	dump, err := httputil.DumpRequestOut(httpReq, true)
	if err != nil {
		return nil, fmt.Errorf("unable to dump http request: %v", err)
	}

	// TODO(cflewis): This might need to be retried if the BCID server is observed to flake out.
	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("unable to contact BCID endpoint: %v", err)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("couldn't read HTTP body: %v", err)
	}
	// Check if the request was blocked by lack of a good token.
	// If so, give a more informative error.
	if resp.StatusCode == http.StatusUnauthorized {
		return nil, fmt.Errorf("GFE blocked access attempt, probably due to bad auth token. Output:\n%s", body)
	}
	// If there was no good status there really isn't anything meaningful to be done at
	// this point. There's no recovery that can be taken.
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("bad HTTP status code %v: body:\n%s\noriginal request:\n%s", resp.StatusCode, body, dump)
	}

	var sResp SCILoResp
	err = json.Unmarshal(body, &sResp)
	if err != nil {
		return nil, fmt.Errorf("unable to marshal HTTP body to SCILoResp: %v", err)
	}
	return &sResp, nil
}

func LUCITokenSource(ctx context.Context) (oauth2.TokenSource, error) {
	authenticator := auth.NewAuthenticator(ctx, auth.SilentLogin, auth.Options{
		Scopes: []string{"https://www.googleapis.com/auth/bcid_verify"},
	})
	return authenticator.TokenSource()
}