package main

import (
	"bytes"
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const maximumBundleBytes = 64 << 20

var objectPrefixPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9-]*/v[1-9][0-9]*$`)

type credentials struct {
	tenancy     string
	user        string
	fingerprint string
	privateKey  *rsa.PrivateKey
}

type stringValues []string

func (values *stringValues) String() string {
	return strings.Join(*values, ",")
}

func (values *stringValues) Set(value string) error {
	*values = append(*values, value)
	return nil
}

type publication struct {
	Prefix          string `json:"prefix"`
	VersionObject   string `json:"version_object"`
	VersionCreated  bool   `json:"version_created"`
	CurrentObject   string `json:"current_object"`
	CurrentPromoted bool   `json:"current_promoted"`
}

type smtpDenySummary struct {
	GeneratedAt          time.Time `json:"generated_at"`
	ExpiresAt            time.Time `json:"expires_at"`
	DeniedIntervalCount  int       `json:"denied_interval_count"`
	DeniedAddressCount   uint64    `json:"denied_address_count"`
	RemovedIntervalCount int       `json:"removed_interval_count"`
}

func readSMTPDenySummary(path string) (*smtpDenySummary, error) {
	if path == "" {
		return nil, nil
	}
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer func() { _ = file.Close() }()
	decoder := json.NewDecoder(io.LimitReader(file, 1<<20))
	decoder.DisallowUnknownFields()
	summary := smtpDenySummary{}
	if err := decoder.Decode(&summary); err != nil {
		return nil, err
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		return nil, errors.New("SMTP deny summary contains trailing data")
	}
	if summary.GeneratedAt.IsZero() ||
		!summary.ExpiresAt.After(summary.GeneratedAt) ||
		summary.DeniedIntervalCount < 0 ||
		summary.RemovedIntervalCount < 0 {
		return nil, errors.New("SMTP deny summary is invalid")
	}
	return &summary, nil
}

func publicationPrefixes(primary string, additional []string) ([]string, error) {
	values := append([]string{primary}, additional...)
	seen := make(map[string]struct{}, len(values))
	result := make([]string, 0, len(values))
	for _, value := range values {
		if !objectPrefixPattern.MatchString(value) {
			return nil, fmt.Errorf("invalid object prefix %q", value)
		}
		if _, exists := seen[value]; exists {
			return nil, fmt.Errorf("duplicate object prefix %q", value)
		}
		seen[value] = struct{}{}
		result = append(result, value)
	}
	return result, nil
}

func publicationObjects(prefix, digest string) (string, string) {
	return prefix + "/versions/" + digest + ".json.gz", prefix + "/current.json.gz"
}

func publishBundle(ctx context.Context, client *http.Client, endpoint, namespace, bucket string, prefixes []string, body []byte, identity credentials, promoteCurrent bool) (string, []publication, error) {
	digest := sha256.Sum256(body)
	digestText := hex.EncodeToString(digest[:])
	publications := make([]publication, 0, len(prefixes))
	for _, prefix := range prefixes {
		versionObject, currentObject := publicationObjects(prefix, digestText)
		created, err := upload(ctx, client, endpoint, namespace, bucket, versionObject, body, identity, true)
		if err != nil {
			return "", nil, err
		}
		if err := verifyObject(ctx, client, endpoint, namespace, bucket, versionObject, body, identity); err != nil {
			return "", nil, err
		}
		if promoteCurrent {
			if _, err := upload(ctx, client, endpoint, namespace, bucket, currentObject, body, identity, false); err != nil {
				return "", nil, err
			}
			if err := verifyObject(ctx, client, endpoint, namespace, bucket, currentObject, body, identity); err != nil {
				return "", nil, err
			}
		}
		publications = append(publications, publication{
			Prefix:          prefix,
			VersionObject:   versionObject,
			VersionCreated:  created,
			CurrentObject:   currentObject,
			CurrentPromoted: promoteCurrent,
		})
	}
	return digestText, publications, nil
}

func writeJSONFile(path string, value any) error {
	contents, err := json.Marshal(value)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return err
	}
	return os.WriteFile(path, append(contents, '\n'), 0o644)
}

func readCredential(directory, name string) (string, error) {
	contents, err := os.ReadFile(filepath.Join(directory, name))
	if err != nil {
		return "", fmt.Errorf("read OCI %s", strings.ReplaceAll(name, "_", " "))
	}
	value := strings.TrimSpace(string(contents))
	if value == "" {
		return "", fmt.Errorf("OCI %s is empty", strings.ReplaceAll(name, "_", " "))
	}
	return value, nil
}

func optionalCredential(directory, name, fallback string) string {
	value, err := readCredential(directory, name)
	if err != nil {
		return fallback
	}
	return value
}

func parsePrivateKey(value string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(value))
	if block == nil {
		return nil, errors.New("decode OCI private key")
	}
	if key, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
		return key, nil
	}
	parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, errors.New("parse OCI private key")
	}
	key, ok := parsed.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("OCI private key is not RSA")
	}
	return key, nil
}

func readCredentials(directory string) (credentials, error) {
	tenancy, err := readCredential(directory, "tenancy_ocid")
	if err != nil {
		return credentials{}, err
	}
	user, err := readCredential(directory, "user_ocid")
	if err != nil {
		return credentials{}, err
	}
	fingerprint, err := readCredential(directory, "fingerprint")
	if err != nil {
		return credentials{}, err
	}
	privateKeyPEM, err := readCredential(directory, "private_key")
	if err != nil {
		return credentials{}, err
	}
	privateKey, err := parsePrivateKey(privateKeyPEM)
	if err != nil {
		return credentials{}, err
	}
	return credentials{tenancy: tenancy, user: user, fingerprint: fingerprint, privateKey: privateKey}, nil
}

func objectURL(endpoint, namespace, bucket, object string) (string, error) {
	parsed, err := url.Parse(endpoint)
	if err != nil {
		return "", err
	}
	if parsed.Scheme != "https" || parsed.Host == "" || parsed.RawQuery != "" || parsed.Fragment != "" {
		return "", errors.New("invalid OCI Object Storage endpoint")
	}
	parsed.Path = "/n/" + namespace + "/b/" + bucket + "/o/" + object
	parsed.RawPath = "/n/" + url.PathEscape(namespace) + "/b/" + url.PathEscape(bucket) + "/o/" + url.PathEscape(object)
	return parsed.String(), nil
}

func signRequest(request *http.Request, body []byte, identity credentials, now time.Time) error {
	digest := sha256.Sum256(body)
	request.Header.Set("Date", now.UTC().Format(http.TimeFormat))
	request.Header.Set("Content-Type", "application/gzip")
	request.Header.Set("Content-Length", strconv.Itoa(len(body)))
	request.Header.Set("X-Content-SHA256", base64.StdEncoding.EncodeToString(digest[:]))
	request.ContentLength = int64(len(body))
	headerNames := "(request-target) host date x-content-sha256 content-type content-length"
	signingString := strings.Join([]string{
		"(request-target): " + strings.ToLower(request.Method) + " " + request.URL.EscapedPath(),
		"host: " + request.URL.Host,
		"date: " + request.Header.Get("Date"),
		"x-content-sha256: " + request.Header.Get("X-Content-SHA256"),
		"content-type: " + request.Header.Get("Content-Type"),
		"content-length: " + request.Header.Get("Content-Length"),
	}, "\n")
	signingDigest := sha256.Sum256([]byte(signingString))
	signature, err := rsa.SignPKCS1v15(rand.Reader, identity.privateKey, crypto.SHA256, signingDigest[:])
	if err != nil {
		return errors.New("sign OCI request")
	}
	keyID := identity.tenancy + "/" + identity.user + "/" + identity.fingerprint
	request.Header.Set("Authorization", fmt.Sprintf(`Signature version="1",keyId="%s",algorithm="rsa-sha256",headers="%s",signature="%s"`, keyID, headerNames, base64.StdEncoding.EncodeToString(signature)))
	return nil
}

func signDownloadRequest(request *http.Request, identity credentials, now time.Time) error {
	request.Header.Set("Date", now.UTC().Format(http.TimeFormat))
	headerNames := "(request-target) host date"
	signingString := strings.Join([]string{
		"(request-target): " + strings.ToLower(request.Method) + " " + request.URL.EscapedPath(),
		"host: " + request.URL.Host,
		"date: " + request.Header.Get("Date"),
	}, "\n")
	signingDigest := sha256.Sum256([]byte(signingString))
	signature, err := rsa.SignPKCS1v15(rand.Reader, identity.privateKey, crypto.SHA256, signingDigest[:])
	if err != nil {
		return errors.New("sign OCI request")
	}
	keyID := identity.tenancy + "/" + identity.user + "/" + identity.fingerprint
	request.Header.Set("Authorization", fmt.Sprintf(`Signature version="1",keyId="%s",algorithm="rsa-sha256",headers="%s",signature="%s"`, keyID, headerNames, base64.StdEncoding.EncodeToString(signature)))
	return nil
}

func upload(ctx context.Context, client *http.Client, endpoint, namespace, bucket, object string, body []byte, identity credentials, immutable bool) (bool, error) {
	target, err := objectURL(endpoint, namespace, bucket, object)
	if err != nil {
		return false, err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPut, target, bytes.NewReader(body))
	if err != nil {
		return false, errors.New("create OCI upload request")
	}
	if immutable {
		request.Header.Set("If-None-Match", "*")
	}
	if err := signRequest(request, body, identity, time.Now()); err != nil {
		return false, err
	}
	response, err := client.Do(request)
	if err != nil {
		return false, errors.New("upload OCI object")
	}
	defer func() { _ = response.Body.Close() }()
	_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 1<<20))
	if immutable && response.StatusCode == http.StatusPreconditionFailed {
		return false, nil
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return false, fmt.Errorf("upload OCI object: HTTP %d", response.StatusCode)
	}
	return true, nil
}

func download(ctx context.Context, client *http.Client, endpoint, namespace, bucket, object string, identity credentials) ([]byte, error) {
	target, err := objectURL(endpoint, namespace, bucket, object)
	if err != nil {
		return nil, err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
	if err != nil {
		return nil, errors.New("create OCI download request")
	}
	if err := signDownloadRequest(request, identity, time.Now()); err != nil {
		return nil, err
	}
	response, err := client.Do(request)
	if err != nil {
		return nil, errors.New("download OCI object")
	}
	defer func() { _ = response.Body.Close() }()
	if response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 1<<20))
		return nil, fmt.Errorf("download OCI object: HTTP %d", response.StatusCode)
	}
	contents, err := io.ReadAll(io.LimitReader(response.Body, maximumBundleBytes+1))
	if err != nil {
		return nil, errors.New("read OCI object")
	}
	if len(contents) > maximumBundleBytes {
		return nil, errors.New("OCI object exceeds maximum bundle size")
	}
	return contents, nil
}

func verifyObject(ctx context.Context, client *http.Client, endpoint, namespace, bucket, object string, expected []byte, identity credentials) error {
	contents, err := download(ctx, client, endpoint, namespace, bucket, object, identity)
	if err != nil {
		return err
	}
	expectedDigest := sha256.Sum256(expected)
	actualDigest := sha256.Sum256(contents)
	if actualDigest != expectedDigest {
		return fmt.Errorf("verify OCI object %q: SHA-256 mismatch", object)
	}
	return nil
}

func main() {
	bundlePath := flag.String("bundle", "/work/output/current.json.gz", "validated admission bundle")
	credentialsDirectory := flag.String("credentials-directory", "/credentials/oci", "OCI credential directory")
	region := flag.String("region", "", "OCI region")
	namespace := flag.String("namespace", "", "OCI Object Storage namespace")
	bucket := flag.String("bucket", "", "OCI Object Storage bucket")
	objectPrefix := flag.String("object-prefix", "admission/v1", "Object Storage prefix for versions and current")
	additionalObjectPrefixes := stringValues{}
	flag.Var(&additionalObjectPrefixes, "additional-object-prefix", "additional Object Storage prefix for the same bundle")
	promoteCurrent := flag.Bool("promote-current", true, "promote the validated bundle after immutable upload")
	endpoint := flag.String("endpoint", "", "OCI Object Storage endpoint")
	resultFile := flag.String("result-file", "", "write sanitized publication result")
	smtpDenySummaryPath := flag.String("smtp-deny-summary", "", "validated SMTP deny summary")
	restoreVersionSHA256 := flag.String("restore-version-sha256", "", "download an immutable version by SHA-256")
	restoreOutput := flag.String("restore-output", "/work/output/restore.json.gz", "write the immutable version for validation")
	flag.Parse()
	*region = optionalCredential(*credentialsDirectory, "region", *region)
	*namespace = optionalCredential(*credentialsDirectory, "object_storage_namespace", *namespace)
	*bucket = optionalCredential(*credentialsDirectory, "bucket_name", *bucket)
	if *region == "" || *namespace == "" || *bucket == "" {
		fmt.Fprintln(os.Stderr, "region, namespace, and bucket are required")
		os.Exit(2)
	}
	if *endpoint == "" {
		*endpoint = "https://objectstorage." + *region + ".oraclecloud.com"
	}
	prefixes, err := publicationPrefixes(*objectPrefix, additionalObjectPrefixes)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
	smtpSummary, err := readSMTPDenySummary(*smtpDenySummaryPath)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
	identity, err := readCredentials(*credentialsDirectory)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()
	client := &http.Client{Timeout: 2 * time.Minute}
	if *restoreVersionSHA256 != "" {
		if !regexp.MustCompile(`^[0-9a-f]{64}$`).MatchString(*restoreVersionSHA256) {
			fmt.Fprintln(os.Stderr, "restore version SHA-256 must be 64 lowercase hexadecimal characters")
			os.Exit(2)
		}
		versionObject, _ := publicationObjects(prefixes[0], *restoreVersionSHA256)
		bundle, err := download(ctx, client, *endpoint, *namespace, *bucket, versionObject, identity)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		digest := sha256.Sum256(bundle)
		if hex.EncodeToString(digest[:]) != *restoreVersionSHA256 {
			fmt.Fprintln(os.Stderr, "downloaded immutable version does not match requested SHA-256")
			os.Exit(1)
		}
		if err := os.MkdirAll(filepath.Dir(*restoreOutput), 0o700); err != nil {
			fmt.Fprintln(os.Stderr, "create restore output directory")
			os.Exit(1)
		}
		if err := os.WriteFile(*restoreOutput, bundle, 0o600); err != nil {
			fmt.Fprintln(os.Stderr, "write restored immutable version")
			os.Exit(1)
		}
		fmt.Printf("restored_bundle_sha256=%s\n", *restoreVersionSHA256)
		return
	}
	bundle, err := os.ReadFile(*bundlePath)
	if err != nil {
		fmt.Fprintln(os.Stderr, "read validated admission bundle")
		os.Exit(1)
	}
	digestText, publications, err := publishBundle(ctx, client, *endpoint, *namespace, *bucket, prefixes, bundle, identity, *promoteCurrent)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	if *resultFile != "" {
		result := struct {
			BundleSHA256    string           `json:"bundle_sha256"`
			Publications    []publication    `json:"publications"`
			SMTPDenySummary *smtpDenySummary `json:"smtp_deny_summary,omitempty"`
		}{BundleSHA256: digestText, Publications: publications, SMTPDenySummary: smtpSummary}
		if err := writeJSONFile(*resultFile, result); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}
	fmt.Printf("bundle_sha256=%s publication_count=%d current_promoted=%t\n", digestText, len(publications), *promoteCurrent)
}