package main

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"
)

const oauthUser = "$oauthtoken"

var challengeRule = regexp.MustCompile(`([A-Za-z]+)="([^"]*)"`)

type quayReader struct {
	api        *url.URL
	registry   string
	client     *http.Client
	candidates string
	mirror     string
}

type manifestDoc struct {
	Manifests json.RawMessage `json:"manifests"`
	Layers    []descriptor    `json:"layers"`
}

type descriptor struct {
	Digest string `json:"digest"`
}

type securityDoc struct {
	Status string        `json:"status"`
	Data   *securityData `json:"data"`
}

type securityData struct {
	Layer     *securityLayer `json:"Layer"`
	Inventory *scanInventory `json:"Inventory"`
}

type securityLayer struct {
	Name     string             `json:"Name"`
	Features *[]securityFeature `json:"Features"`
}

type securityFeature struct {
	Name            string          `json:"Name"`
	Version         string          `json:"Version"`
	AddedBy         string          `json:"AddedBy"`
	Vulnerabilities []vulnerability `json:"Vulnerabilities"`
}

type vulnerability struct {
	Severity string `json:"Severity"`
}

func loadRegistry() (*quayReader, error) {
	registry := os.Getenv("REGISTRY")
	if registry == "" {
		return nil, errors.New("REGISTRY is required")
	}
	apiValue := os.Getenv("QUAY_API")
	if apiValue == "" {
		return nil, errors.New("QUAY_API is required")
	}
	api, err := url.Parse(apiValue)
	if err != nil || api.Host == "" || (api.Scheme != "http" && api.Scheme != "https") {
		return nil, errors.New("QUAY_API must be an HTTP or HTTPS origin")
	}
	transport := http.DefaultTransport.(*http.Transport).Clone()
	if api.Scheme == "https" {
		roots, err := x509.SystemCertPool()
		if err != nil {
			return nil, fmt.Errorf("load system roots: %w", err)
		}
		certificate, err := os.ReadFile(valueOrDefault("REGISTRY_CA", "/registry-ca/ca.crt"))
		if err != nil {
			return nil, fmt.Errorf("read registry CA: %w", err)
		}
		if !roots.AppendCertsFromPEM(certificate) {
			return nil, errors.New("registry CA contains no certificate")
		}
		transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}
	}
	candidates, err := readSecret(valueOrDefault("CANDIDATES_TOKEN", "/candidates-reader/token"))
	if err != nil {
		return nil, err
	}
	mirror, err := readSecret(valueOrDefault("MIRROR_TOKEN", "/mirror-reader/token"))
	if err != nil {
		return nil, err
	}
	return &quayReader{
		api:        api,
		registry:   registry,
		client:     &http.Client{Timeout: 20 * time.Second, Transport: transport},
		candidates: candidates,
		mirror:     mirror,
	}, nil
}

func readSecret(path string) (string, error) {
	content, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("read %s: %w", path, err)
	}
	value := strings.TrimSpace(string(content))
	if value == "" {
		return "", fmt.Errorf("secret %s is empty", path)
	}
	return value, nil
}

func (q *quayReader) scan(ctx context.Context, repository, digest, threshold string) (string, error) {
	content, found, err := q.manifest(ctx, repository, digest)
	if err != nil {
		return "", err
	}
	if !found {
		return scanPending, nil
	}
	if fmt.Sprintf("sha256:%x", sha256.Sum256(content)) != digest {
		return "", errors.New("candidate digest does not match manifest")
	}
	digests, err := scanDigests(content, digest)
	if err != nil {
		return "", err
	}
	limit, err := severityRank(threshold)
	if err != nil {
		return "", err
	}
	for _, child := range digests {
		result, found, err := q.security(ctx, repository, child)
		if err != nil {
			return "", err
		}
		if !found || result.Status == "queued" || result.Status == "scanning" || result.Status == "unsupported" {
			return scanPending, nil
		}
		if result.Status != "scanned" {
			return "", fmt.Errorf("clair returned unusable status %q for %s", result.Status, child)
		}
		if result.Data == nil || result.Data.Layer == nil || result.Data.Layer.Features == nil {
			return "", fmt.Errorf("clair scan data is unavailable for %s", child)
		}
		manifest := content
		if child != digest {
			manifest, found, err = q.manifest(ctx, repository, child)
			if err != nil || !found {
				return scanPending, err
			}
		}
		if fmt.Sprintf("sha256:%x", sha256.Sum256(manifest)) != child {
			return "", errors.New("child digest does not match manifest")
		}
		var document manifestDoc
		if err := json.Unmarshal(manifest, &document); err != nil {
			return "", errors.New("candidate manifest is invalid")
		}
		if len(document.Manifests) != 0 || len(document.Layers) == 0 {
			return scanRejected, nil
		}
		layers := map[string]bool{}
		for _, layer := range document.Layers {
			if !digestRule.MatchString(layer.Digest) {
				return "", errors.New("candidate layer digest is invalid")
			}
			layers[layer.Digest] = true
		}
		if err := checkInventory(result.Data, child, layers); err != nil {
			return scanRejected, nil
		}
		for _, feature := range *result.Data.Layer.Features {
			for _, finding := range feature.Vulnerabilities {
				rank, err := severityRank(finding.Severity)
				if err != nil {
					return "", err
				}
				if rank == 0 || rank >= limit {
					return scanRejected, nil
				}
			}
		}
	}
	return scanAccepted, nil
}

func scanDigests(content []byte, digest string) ([]string, error) {
	var manifest manifestDoc
	if err := json.Unmarshal(content, &manifest); err != nil {
		return nil, fmt.Errorf("decode candidate manifest: %w", err)
	}
	if len(manifest.Manifests) == 0 {
		return []string{digest}, nil
	}
	var children []descriptor
	if err := json.Unmarshal(manifest.Manifests, &children); err != nil {
		return nil, fmt.Errorf("decode candidate index: %w", err)
	}
	if len(children) == 0 {
		return nil, errors.New("candidate index contains no manifests")
	}
	digests := make([]string, 0, len(children))
	for _, child := range children {
		if !digestRule.MatchString(child.Digest) {
			return nil, fmt.Errorf("candidate index contains invalid digest %q", child.Digest)
		}
		digests = append(digests, child.Digest)
	}
	return digests, nil
}

func severityRank(value string) (int, error) {
	switch strings.ToLower(value) {
	case "unknown":
		return 0, nil
	case "negligible":
		return 1, nil
	case "low":
		return 2, nil
	case "medium":
		return 3, nil
	case "high":
		return 4, nil
	case "critical":
		return 5, nil
	default:
		return 0, fmt.Errorf("unknown Clair severity %q", value)
	}
}

func (q *quayReader) exists(ctx context.Context, repository, digest string) (bool, error) {
	_, found, err := q.manifest(ctx, repository, digest)
	return found, err
}

func (q *quayReader) physicallyPresent(ctx context.Context, repository, digest string) (bool, error) {
	response, err := q.apiGet(ctx, repository, digest, "/security?vulnerabilities=false")
	if err != nil {
		return false, err
	}
	defer func() { _ = response.Body.Close() }()
	switch response.StatusCode {
	case http.StatusOK:
		return true, nil
	case http.StatusNotFound:
		return false, nil
	default:
		return false, fmt.Errorf("quay physical manifest lookup failed with HTTP %d", response.StatusCode)
	}
}

func (q *quayReader) security(ctx context.Context, repository, digest string) (securityDoc, bool, error) {
	response, err := q.apiGet(ctx, repository, digest, "/security?vulnerabilities=true")
	if err != nil {
		return securityDoc{}, false, err
	}
	defer func() { _ = response.Body.Close() }()
	if response.StatusCode == http.StatusNotFound {
		return securityDoc{}, false, nil
	}
	if response.StatusCode != http.StatusOK {
		return securityDoc{}, false, fmt.Errorf("quay security lookup failed with HTTP %d", response.StatusCode)
	}
	var result securityDoc
	if err := json.NewDecoder(io.LimitReader(response.Body, 8<<20)).Decode(&result); err != nil {
		return securityDoc{}, false, fmt.Errorf("decode Quay security result: %w", err)
	}
	return result, true, nil
}

func (q *quayReader) apiGet(ctx context.Context, repository, digest, suffix string) (*http.Response, error) {
	endpoint := *q.api
	endpoint.Path = strings.TrimRight(endpoint.Path, "/") + "/api/v1/repository/" + repository + "/manifest/" + digest
	endpoint.RawQuery = ""
	if queryIndex := strings.IndexByte(suffix, '?'); queryIndex >= 0 {
		endpoint.Path += suffix[:queryIndex]
		endpoint.RawQuery = suffix[queryIndex+1:]
	} else {
		endpoint.Path += suffix
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.String(), nil)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Authorization", "Bearer "+q.token(repository))
	return q.client.Do(request)
}

func (q *quayReader) manifest(ctx context.Context, repository, digest string) ([]byte, bool, error) {
	endpoint := fmt.Sprintf("%s://%s/v2/%s/manifests/%s", q.api.Scheme, q.registry, repository, digest)
	send := func(authorization string) (*http.Response, error) {
		request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
		if err != nil {
			return nil, err
		}
		request.Header.Set("Accept", strings.Join([]string{
			"application/vnd.oci.image.index.v1+json",
			"application/vnd.oci.image.manifest.v1+json",
			"application/vnd.docker.distribution.manifest.list.v2+json",
			"application/vnd.docker.distribution.manifest.v2+json",
		}, ", "))
		if authorization != "" {
			request.Header.Set("Authorization", authorization)
		}
		return q.client.Do(request)
	}
	response, err := send("")
	if err != nil {
		return nil, false, err
	}
	if response.StatusCode == http.StatusUnauthorized {
		challenge := response.Header.Get("WWW-Authenticate")
		if err := response.Body.Close(); err != nil {
			return nil, false, fmt.Errorf("close registry challenge: %w", err)
		}
		token, err := q.bearerToken(ctx, challenge, repository)
		if err != nil {
			return nil, false, err
		}
		response, err = send("Bearer " + token)
		if err != nil {
			return nil, false, err
		}
	}
	defer func() { _ = response.Body.Close() }()
	if response.StatusCode == http.StatusNotFound {
		return nil, false, nil
	}
	if response.StatusCode != http.StatusOK {
		return nil, false, fmt.Errorf("registry manifest lookup failed with HTTP %d", response.StatusCode)
	}
	content, err := io.ReadAll(io.LimitReader(response.Body, 8<<20))
	if err != nil {
		return nil, false, fmt.Errorf("read candidate manifest: %w", err)
	}
	return content, true, nil
}

func (q *quayReader) bearerToken(ctx context.Context, challenge, repository string) (string, error) {
	values := make(map[string]string)
	for _, match := range challengeRule.FindAllStringSubmatch(challenge, -1) {
		values[strings.ToLower(match[1])] = match[2]
	}
	realm, err := url.Parse(values["realm"])
	if err != nil || realm.Host == "" {
		return "", errors.New("registry challenge omitted a valid realm")
	}
	query := realm.Query()
	query.Set("service", values["service"])
	query.Set("scope", values["scope"])
	query.Set("account", oauthUser)
	realm.RawQuery = query.Encode()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, realm.String(), nil)
	if err != nil {
		return "", err
	}
	request.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(oauthUser+":"+q.token(repository))))
	response, err := q.client.Do(request)
	if err != nil {
		return "", err
	}
	defer func() { _ = response.Body.Close() }()
	if response.StatusCode != http.StatusOK {
		return "", fmt.Errorf("registry authentication failed with HTTP %d", response.StatusCode)
	}
	var result struct {
		Token       string `json:"token"`
		AccessToken string `json:"access_token"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&result); err != nil {
		return "", fmt.Errorf("decode registry token: %w", err)
	}
	if result.Token != "" {
		return result.Token, nil
	}
	if result.AccessToken != "" {
		return result.AccessToken, nil
	}
	return "", errors.New("registry authentication returned no token")
}

func (q *quayReader) token(repository string) string {
	if strings.HasPrefix(repository, "candidates/") {
		return q.candidates
	}
	return q.mirror
}