// Copyright 2025 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 cryptoutil

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"testing"

	"github.com/secure-systems-lab/go-securesystemslib/encrypted"
	"github.com/stretchr/testify/require"
)

// pemsigstorekey copied from sigstore/cosign tests; passphrase is "hello".
const pemsigstorekey = `-----BEGIN ENCRYPTED SIGSTORE PRIVATE KEY-----
eyJrZGYiOnsibmFtZSI6InNjcnlwdCIsInBhcmFtcyI6eyJOIjozMjc2OCwiciI6
OCwicCI6MX0sInNhbHQiOiI3T3VGd2VsbWZZNXVId2NoaURSc210anNwZ2ZlZjFG
Mk5lOGFDTjVLYVpZPSJ9LCJjaXBoZXIiOnsibmFtZSI6Im5hY2wvc2VjcmV0Ym94
Iiwibm9uY2UiOiJQNHk4OGhCb3ZTa09MbXN0bFVBaGJwdDJ0K2xTNUxQSCJ9LCJj
aXBoZXJ0ZXh0IjoiMnB1QzdyZldJOWh3bnJlQ2s4aUZDRlVwQlRrSzRJNlIvbFBF
cnBDekpXUGpJWXl4eGVIL1A2VW52cFJHdVhla1NNb3JMdGhLamdoQ1JlNy82NDVH
QWtoVm1LRC92eEF0S2EvbE1abENSQ3FlekJGUFd1dzNpeFRtZ2xhb2J1ZFVSbUVs
bmNGOGlZbzBTMVl6Y1ZOMVFwY2J2c0dNcUlYRzVlbmdteGp5dCtBcXlyZTF0Q0Y0
V01tU1BlaEljNlBqd2h1Q2xHaVpJUWRvTGc9PSJ9
-----END ENCRYPTED SIGSTORE PRIVATE KEY-----`

// pemcosignkey copied from sigstore/cosign tests; passphrase is "hello".
const pemcosignkey = `-----BEGIN ENCRYPTED COSIGN PRIVATE KEY-----
eyJrZGYiOnsibmFtZSI6InNjcnlwdCIsInBhcmFtcyI6eyJOIjozMjc2OCwiciI6
OCwicCI6MX0sInNhbHQiOiJ4WWdoc09JTUxUWGNOT0RsclNIOUNKc1FlOVFnZmN1
cmUrMXlLdHh1TlkwPSJ9LCJjaXBoZXIiOnsibmFtZSI6Im5hY2wvc2VjcmV0Ym94
Iiwibm9uY2UiOiI0cS9PSlVmaXJkSUkrUjZ0ajZBMmcyQ0JqL25xdFNicCJ9LCJj
aXBoZXJ0ZXh0IjoiKzB4Q3NzcFN0WStBczdKanJpOWtsbHBWd2JhcUI4ZWJNdWto
eS9aVE1MSXRsL3B1YS9jWVJvbytLRGxMWWdmOW1kSjk4K1FnQW9oTktoYnJPMTcw
MHdBY1JTMjFDOE4zQUNJRUVZaWpOMllBNnMraGJSbkhjUnd4eGhDMDFtb2FvL0dO
Y1pmbEJheXZMV3pXblo4d2NDZ2ZpT1o1VXlRTEFJMHh0dnR6dEh3cTdDV1Vhd3V4
RlhlNDZzck9TUE9SNHN6bytabWErUGovSFE9PSJ9
-----END ENCRYPTED COSIGN PRIVATE KEY-----`

func TestTryParseKeyFromReaderWithPassword_SigstoreEncrypted_Success(t *testing.T) {
	parsed, err := TryParseKeyFromReaderWithPassword(bytes.NewReader([]byte(pemsigstorekey)), []byte("hello"))
	require.NoError(t, err)

	switch parsed.(type) {
	case *rsa.PrivateKey, *ecdsa.PrivateKey, ed25519.PrivateKey:
		// ok
	default:
		t.Fatalf("unexpected key type %T", parsed)
	}
}

func TestTryParseKeyFromReaderWithPassword_SigstoreEncrypted_WrongPass(t *testing.T) {
	_, err := TryParseKeyFromReaderWithPassword(bytes.NewReader([]byte(pemsigstorekey)), []byte("wrong"))
	require.Error(t, err)
}

func TestTryParseKeyFromReaderWithPassword_SigstoreEncrypted_NoPass(t *testing.T) {
	_, err := TryParseKeyFromReaderWithPassword(bytes.NewReader([]byte(pemsigstorekey)), nil)
	require.Error(t, err)
}

func TestTryParseKeyFromReaderWithPassword_CosignEncrypted_Success(t *testing.T) {
	parsed, err := TryParseKeyFromReaderWithPassword(bytes.NewReader([]byte(pemcosignkey)), []byte("hello"))
	require.NoError(t, err)

	switch parsed.(type) {
	case *rsa.PrivateKey, *ecdsa.PrivateKey, ed25519.PrivateKey:
		// ok
	default:
		t.Fatalf("unexpected key type %T", parsed)
	}
}

func TestTryParseKeyFromReaderWithPassword_CosignEncrypted_WrongPass(t *testing.T) {
	_, err := TryParseKeyFromReaderWithPassword(bytes.NewReader([]byte(pemcosignkey)), []byte("wrong"))
	require.Error(t, err)
}

func TestTryParseKeyFromReaderWithPassword_CosignEncrypted_NoPass(t *testing.T) {
	_, err := TryParseKeyFromReaderWithPassword(bytes.NewReader([]byte(pemcosignkey)), nil)
	require.Error(t, err)
}

func TestTryParseKeyFromReaderWithPassword_SigstoreEmptyPassphrase(t *testing.T) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	require.NoError(t, err)

	der, err := x509.MarshalPKCS8PrivateKey(priv)
	require.NoError(t, err)

	encBytes, err := encrypted.Encrypt(der, []byte(""))
	require.NoError(t, err)

	pemBytes := pem.EncodeToMemory(&pem.Block{
		Type:  "ENCRYPTED SIGSTORE PRIVATE KEY",
		Bytes: encBytes,
	})

	parsed, err := TryParseKeyFromReaderWithPassword(bytes.NewReader(pemBytes), []byte(""))
	require.NoError(t, err)

	if _, ok := parsed.(*ecdsa.PrivateKey); !ok {
		t.Fatalf("expected *ecdsa.PrivateKey, got %T", parsed)
	}
}

func TestTryParseKeyFromReaderWithPassword_CosignEmptyPassphrase(t *testing.T) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	require.NoError(t, err)

	der, err := x509.MarshalPKCS8PrivateKey(priv)
	require.NoError(t, err)

	encBytes, err := encrypted.Encrypt(der, []byte(""))
	require.NoError(t, err)

	pemBytes := pem.EncodeToMemory(&pem.Block{
		Type:  "ENCRYPTED COSIGN PRIVATE KEY",
		Bytes: encBytes,
	})

	parsed, err := TryParseKeyFromReaderWithPassword(bytes.NewReader(pemBytes), []byte(""))
	require.NoError(t, err)

	if _, ok := parsed.(*ecdsa.PrivateKey); !ok {
		t.Fatalf("expected *ecdsa.PrivateKey, got %T", parsed)
	}
}

func TestTryParseKeyFromReaderWithPassword_CosignEmptyPassphraseBytes(t *testing.T) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	require.NoError(t, err)

	der, err := x509.MarshalPKCS8PrivateKey(priv)
	require.NoError(t, err)

	encBytes, err := encrypted.Encrypt(der, []byte(""))
	require.NoError(t, err)

	pemBytes := pem.EncodeToMemory(&pem.Block{
		Type:  "ENCRYPTED COSIGN PRIVATE KEY",
		Bytes: encBytes,
	})

	parsed, err := TryParseKeyFromReaderWithPassword(bytes.NewReader(pemBytes), []byte{})
	require.NoError(t, err)

	if _, ok := parsed.(*ecdsa.PrivateKey); !ok {
		t.Fatalf("expected *ecdsa.PrivateKey, got %T", parsed)
	}
}
