SignEncrypt.cs

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Linq;
using System.Security.Cryptography.Xml;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Forms;
using System.Xml.Linq;
using Net.Pkcs11Interop.Common;
using Net.Pkcs11Interop.HighLevelAPI;

namespace CryptokiTokenBrowser
{
    public partial class SignEncrypt : Form
    {
        public ISession? session { get; set; } = null;
        public IObjectHandle? encryptKey { get; set; } = null;
        public IObjectHandle? signKey { get; set; } = null;
        public IMechanism? encryptMechanism { get; set; } = null;
        public IMechanism? signMechanism { get; set; } = null;
        public ISlot? selectedSlot = null;

        private Dictionary<string, IObjectHandle> labelToKeyMap =
            new Dictionary<string, IObjectHandle>();

        public SignEncrypt()
        {
            InitializeComponent();
        }

        private (byte[] signature, byte[] encryptedData) SignEncryptProcess(
            byte[] data,
            ISession? session,
            IObjectHandle? encryptKey,
            IObjectHandle? signKey,
            IMechanism? encryptMechanism,
            IMechanism? signMechanism
        )
        {
            byte[] signature = [];
            byte[] encryptedData = [];

            if (data == null || data.Length == 0)
            {
                MessageBox.Show(
                    "Data to encrypt cannot be null or empty.",
                    "Input Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Warning
                );
                return ([], []);
            }

            if (session == null)
            {
                MessageBox.Show(
                    "Session is not initialized.",
                    "Session Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return ([], []);
            }

            if (signKey == null)
            {
                MessageBox.Show(
                    "signKey is not initialized.",
                    "signKey Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return ([], []);
            }
            if (encryptKey == null)
            {
                MessageBox.Show(
                    "encryptKey is not initialized.",
                    "encryptKey Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return ([], []);
            }

            if (encryptMechanism == null)
            {
                MessageBox.Show(
                    "encryptMechanism is not initialized.",
                    "encryptMechanism Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return ([], []);
            }

            if (signMechanism == null)
            {
                MessageBox.Show(
                    "signMechanism is not initialized.",
                    "signMechanism Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return ([], []);
            }

            session.SignEncrypt(
                signMechanism,
                signKey,
                encryptMechanism,
                encryptKey,
                data,
                out signature,
                out encryptedData
            );

            return (signature, encryptedData);
        }

        private void PopulateEncryptKeyComboBox()
        {
            CbEncryptKey.Items.Clear();
            labelToKeyMap.Clear();

            if (session == null)
            {
                MessageBox.Show(
                    "Session is not initialized.",
                    "Session Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return;
            }

            List<IObjectHandle> secretKeyObjects = session.FindAllObjects(
                new List<IObjectAttribute>
                {
                    session.Factories.ObjectAttributeFactory.Create(
                        CKA.CKA_CLASS,
                        CKO.CKO_SECRET_KEY
                    ),
                }
            );

            List<IObjectHandle> privKeyObjects = session.FindAllObjects(
                new List<IObjectAttribute>
                {
                    session.Factories.ObjectAttributeFactory.Create(
                        CKA.CKA_CLASS,
                        CKO.CKO_PRIVATE_KEY
                    ),
                }
            );

            List<IObjectHandle> pubKeyObjects = session.FindAllObjects(
                new List<IObjectAttribute>
                {
                    session.Factories.ObjectAttributeFactory.Create(
                        CKA.CKA_CLASS,
                        CKO.CKO_PUBLIC_KEY
                    ),
                }
            );

            var encryptObjects = secretKeyObjects.Union(pubKeyObjects).ToList();
            var signObject = secretKeyObjects.Union(privKeyObjects).ToList();

            foreach (IObjectHandle keyObject in encryptObjects)
            {
                List<IObjectAttribute> attributes = session.GetAttributeValue(
                    keyObject,
                    new List<CKA> { CKA.CKA_LABEL }
                );

                string label =
                    attributes
                        .FirstOrDefault(a => a.Type == (ulong)CKA.CKA_LABEL)
                        ?.GetValueAsString() ?? "(No Label)";

                labelToKeyMap[label] = keyObject;
                CbEncryptKey.Items.Add(label);
            }
            foreach (IObjectHandle signature in signObject)
            {
                List<IObjectAttribute> attributes = session.GetAttributeValue(
                    signature,
                    new List<CKA> { CKA.CKA_LABEL }
                );

                string label =
                    attributes
                        .FirstOrDefault(a => a.Type == (ulong)CKA.CKA_LABEL)
                        ?.GetValueAsString() ?? "(No Label)";

                labelToKeyMap[label] = signature;
                CbSignKey.Items.Add(label);
            }
        }

        private void PopulateMechanismComboBox()
        {
            if (selectedSlot == null)
            {
                MessageBox.Show(
                    "Session is not initialized.",
                    "Session Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return;
            }

            List<CKM> mechanisms = selectedSlot.GetMechanismList();
            CbMechanismEnc.Items.AddRange(mechanisms.Cast<object>().ToArray());
            CbMechanismSign.Items.AddRange(mechanisms.Cast<object>().ToArray());
        }

        private void SignEncrypt_Load(object sender, EventArgs e)
        {
            if (session == null)
            {
                MessageBox.Show(
                    "Session is not initialized.",
                    "Session Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
                return;
            }
            PopulateMechanismComboBox();
            PopulateEncryptKeyComboBox();
        }

        private void BtnInitEnc_Click(object sender, EventArgs e)
        {
            try
            {
                byte[] ivBytes = new byte[16];
                string iv = TxtParamEnc.Text;
                ivBytes = ConvertUtils.HexStringToBytes(iv);
                string? selectedEncryptKeyLabel = CbEncryptKey.SelectedItem?.ToString();
                string? selectedEncryptMechanismString = CbMechanismEnc.SelectedItem?.ToString();
                CKM mechanismType;

                if (
                    selectedEncryptKeyLabel != null
                    && labelToKeyMap.TryGetValue(
                        selectedEncryptKeyLabel,
                        out IObjectHandle? selectedKeyObject
                    )
                )
                {
                    encryptKey = selectedKeyObject;
                }
                else
                {
                    MessageBox.Show("Selected encryption key not found.");
                    return;
                }

                if (
                    !string.IsNullOrEmpty(selectedEncryptMechanismString)
                    && Enum.TryParse<CKM>(selectedEncryptMechanismString, out mechanismType)
                    && string.IsNullOrEmpty(iv)
                )
                {
                    encryptMechanism = session?.Factories.MechanismFactory.Create(mechanismType);
                }
                else if (
                    !string.IsNullOrEmpty(selectedEncryptMechanismString)
                    && Enum.TryParse<CKM>(selectedEncryptMechanismString, out mechanismType)
                    && !string.IsNullOrEmpty(iv)
                )
                {
                    encryptMechanism = session?.Factories.MechanismFactory.Create(
                        mechanismType,
                        ivBytes
                    );
                }
                else
                {
                    MessageBox.Show("Invalid encryption mechanism selected.");
                    return;
                }
            }
            catch (Exception ex)
            {
                MessageBox.Show(
                    "Error setting encryption key/mechanism: " + ex.Message,
                    "Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
            }
        }

        private void BtnInitSign_Click(object sender, EventArgs e)
        {
            try
            {
                try
                {
                    string? selectedSignKeyLabel = CbSignKey.SelectedItem?.ToString();
                    string? selectedSignMechanismString = CbMechanismSign.SelectedItem?.ToString();

                    if (
                        selectedSignKeyLabel != null
                        && labelToKeyMap.TryGetValue(
                            selectedSignKeyLabel,
                            out IObjectHandle? selectedKeyObject
                        )
                    )
                    {
                        signKey = selectedKeyObject;
                    }
                    else
                    {
                        MessageBox.Show("Selected encryption key not found.");
                        return;
                    }

                    if (
                        !string.IsNullOrEmpty(selectedSignMechanismString)
                        && Enum.TryParse<CKM>(selectedSignMechanismString, out CKM mechanismType)
                    )
                    {
                        signMechanism = session?.Factories.MechanismFactory.Create(mechanismType);
                    }
                    else
                    {
                        MessageBox.Show("Invalid encryption mechanism selected.");
                        return;
                    }
                }
                catch (Exception ex)
                {
                    MessageBox.Show(
                        "Error setting encryption key/mechanism: " + ex.Message,
                        "Error",
                        MessageBoxButtons.OK,
                        MessageBoxIcon.Error
                    );
                }
            }
            catch (Exception ex)
            {
                MessageBox.Show(
                    "An error occurred while initializing signature: " + ex.Message,
                    "Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
            }
        }

        private void BtnSignEnc_Click(object sender, EventArgs e)
        {
            try
            {
                string data = string.Empty;
                string signatureString = string.Empty;
                string encryptedString = string.Empty;
                data = TxtInput.Text;
                byte[] dataBytes = [];

                if (RadioInputAsc.Checked)
                {
                    dataBytes = ConvertUtils.Utf8StringToBytes(data);
                    //int paddingLength = 16 - (dataBytes.Length % 16);
                    //if (paddingLength == 16)
                    //    paddingLength = 0;

                    //byte[] paddedData = new byte[dataBytes.Length + paddingLength];
                    //Array.Copy(dataBytes, paddedData, dataBytes.Length);

                    //dataBytes = paddedData;
                }
                else if (RadioInputHex.Checked)
                {
                    //dataBytes = ConvertUtils.HexStringToBytes(data);
                    //int paddingLength = 16 - (dataBytes.Length % 16);
                    //if (paddingLength == 16)
                    //    paddingLength = 0;

                    //byte[] paddedData = new byte[dataBytes.Length + paddingLength];
                    //Array.Copy(dataBytes, paddedData, dataBytes.Length);

                    //dataBytes = paddedData;
                }
                else
                {
                    MessageBox.Show(
                        "Please select an input format (ASCII or Hex).",
                        "Input Format Error",
                        MessageBoxButtons.OK,
                        MessageBoxIcon.Warning
                    );
                    return;
                }

                var (signature, encryptedData) = SignEncryptProcess(
                    dataBytes,
                    session,
                    encryptKey,
                    signKey,
                    encryptMechanism,
                    signMechanism
                );
                signatureString = ConvertUtils.BytesToHexString(signature);
                encryptedString = ConvertUtils.BytesToHexString(encryptedData);

                TxtSignature.Text = signatureString;
                TxtOutput.Text = encryptedString;
            }
            catch (Exception ex)
            {
                MessageBox.Show(
                    "An error occurred while signing and encrypting: " + ex.Message,
                    "Error",
                    MessageBoxButtons.OK,
                    MessageBoxIcon.Error
                );
            }
        }

        private void BtnCancel_Click(object sender, EventArgs e)
        {
            this.Close();
        }
    }
}