//go:build linux

package executor

import (
	"bytes"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"io"
	"os"
	"path/filepath"
	"syscall"

	"github.com/containerd/continuity/fs"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/util/openfile"
	"github.com/pkg/errors"
)

var linuxSystemCertFiles = []string{
	"/etc/ssl/certs/ca-certificates.crt",
	"/etc/pki/tls/certs/ca-bundle.crt",
	"/etc/ssl/ca-bundle.pem",
	"/etc/pki/tls/cacert.pem",
	"/etc/pki/ca-trust/extracted/pem/tls-ca-bundle.pem",
	"/etc/ssl/cert.pem",
}

var (
	proxyCABegin = []byte("\n# buildkit proxy CA begin\n")
	proxyCAEnd   = []byte("# buildkit proxy CA end\n")
)

const maxCertBundleBytes = 10 << 20

// InjectProxyCA appends caPEM to the rootfs trust bundle used by common Linux
// TLS stacks and returns a cleanup that removes only the injected CA.
func InjectProxyCA(rootfsPath string, caPEM []byte) (func() error, error) {
	if len(caPEM) == 0 {
		return func() error { return nil }, nil
	}
	cert, err := firstCertificate(caPEM)
	if err != nil {
		return nil, err
	}
	certSum := sha256.Sum256(cert.Raw)

	var bundle string
	for _, name := range linuxSystemCertFiles {
		p, err := fs.RootPath(rootfsPath, name)
		if err != nil {
			return nil, errors.Wrapf(err, "failed to resolve certificate bundle %s", name)
		}
		if st, err := os.Stat(p); err == nil && st.Mode().IsRegular() {
			bundle = p
			break
		}
	}
	if bundle == "" {
		return func() error { return nil }, nil
	}
	bundleRel, err := filepath.Rel(rootfsPath, bundle)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to relativize certificate bundle %s", bundle)
	}
	root, err := os.OpenRoot(rootfsPath)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	original, st, err := readCertBundle(root, bundleRel)
	if err != nil {
		_ = root.Close()
		return nil, err
	}
	if containsCertificate(original, certSum) {
		if err := root.Close(); err != nil {
			return nil, errors.WithStack(err)
		}
		return func() error { return nil }, nil
	}
	next := append([]byte{}, original...)
	next = append(next, proxyCABegin...)
	next = append(next, caPEM...)
	if len(next) > 0 && next[len(next)-1] != '\n' {
		next = append(next, '\n')
	}
	next = append(next, proxyCAEnd...)
	if err := writeCertBundle(root, bundleRel, next, st); err != nil {
		_ = root.Close()
		return nil, err
	}

	return func() (retErr error) {
		defer func() {
			if err := root.Close(); retErr == nil && err != nil {
				retErr = errors.WithStack(err)
			}
		}()

		current, st, err := readCertBundle(root, bundleRel)
		if err != nil {
			if errors.Is(err, os.ErrNotExist) {
				return nil
			}
			return err
		}
		cleaned := removeInjectedCA(current, certSum)
		if bytes.Equal(current, cleaned) {
			return nil
		}
		return writeCertBundle(root, bundleRel, cleaned, st)
	}, nil
}

func readCertBundle(root *os.Root, path string) ([]byte, os.FileInfo, error) {
	f, err := openfile.RegularInRoot(root.Name(), path)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}
	defer f.Close()

	st, err := f.Stat()
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}

	limited := &io.LimitedReader{R: f, N: maxCertBundleBytes + 1}
	dt, err := io.ReadAll(limited)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}
	if limited.N == 0 {
		return nil, nil, errors.Errorf("%s exceeds %d bytes", path, maxCertBundleBytes)
	}
	return dt, st, nil
}

func firstCertificate(dt []byte) (*x509.Certificate, error) {
	for {
		block, rest := pem.Decode(dt)
		if block == nil {
			return nil, errors.New("proxy CA PEM does not contain a certificate")
		}
		dt = rest
		if block.Type != "CERTIFICATE" {
			continue
		}
		cert, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return nil, errors.WithStack(err)
		}
		return cert, nil
	}
}

func containsCertificate(dt []byte, sum [sha256.Size]byte) bool {
	for {
		block, rest := pem.Decode(dt)
		if block == nil {
			return false
		}
		dt = rest
		if block.Type != "CERTIFICATE" {
			continue
		}
		if sha256.Sum256(block.Bytes) == sum {
			return true
		}
	}
}

func removeInjectedCA(dt []byte, sum [sha256.Size]byte) []byte {
	begin := bytes.Index(dt, proxyCABegin)
	if begin >= 0 {
		end := bytes.Index(dt[begin+len(proxyCABegin):], proxyCAEnd)
		if end >= 0 {
			end += begin + len(proxyCABegin) + len(proxyCAEnd)
			block := dt[begin:end]
			if containsCertificate(block, sum) {
				out := append([]byte{}, dt[:begin]...)
				out = append(out, dt[end:]...)
				return out
			}
		}
	}

	var out []byte
	for len(dt) > 0 {
		idx := bytes.Index(dt, []byte("-----BEGIN "))
		if idx < 0 {
			out = append(out, dt...)
			break
		}
		out = append(out, dt[:idx]...)
		block, rest := pem.Decode(dt[idx:])
		if block == nil {
			out = append(out, dt[idx:]...)
			break
		}
		consumed := len(dt[idx:]) - len(rest)
		if block.Type != "CERTIFICATE" || sha256.Sum256(block.Bytes) != sum {
			out = append(out, dt[idx:idx+consumed]...)
		}
		dt = rest
	}
	return out
}

func writeCertBundle(root *os.Root, path string, dt []byte, st os.FileInfo) error {
	tmpName := filepath.Join(filepath.Dir(path), ".buildkit-ca-"+identity.NewID())
	tmp, err := root.OpenFile(tmpName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return errors.WithStack(err)
	}
	defer root.Remove(tmpName)
	if _, err := tmp.Write(dt); err != nil {
		tmp.Close()
		return errors.WithStack(err)
	}
	if err := tmp.Chmod(st.Mode()); err != nil {
		tmp.Close()
		return errors.WithStack(err)
	}
	if sys, ok := st.Sys().(*syscall.Stat_t); ok {
		if err := tmp.Chown(int(sys.Uid), int(sys.Gid)); err != nil {
			tmp.Close()
			return errors.WithStack(err)
		}
	}
	if err := tmp.Close(); err != nil {
		return errors.WithStack(err)
	}
	return errors.WithStack(root.Rename(tmpName, path))
}