using System.Security.Cryptography;

namespace OF_DL.Crypto;

public class Padding
{
    public static byte[] AddPKCS7Padding(byte[] data, int k)
    {
        int m = k - data.Length % k;

        byte[] padding = new byte[m];
        Array.Fill(padding, (byte)m);

        byte[] paddedBytes = new byte[data.Length + padding.Length];
        Buffer.BlockCopy(data, 0, paddedBytes, 0, data.Length);
        Buffer.BlockCopy(padding, 0, paddedBytes, data.Length, padding.Length);

        return paddedBytes;
    }

    public static byte[] RemovePKCS7Padding(byte[] paddedByteArray)
    {
        byte last = paddedByteArray[^1];
        if (paddedByteArray.Length <= last)
        {
            return paddedByteArray;
        }

        return SubArray(paddedByteArray, 0, paddedByteArray.Length - last);
    }

    public static T[] SubArray<T>(T[] arr, int start, int length)
    {
        T[] result = new T[length];
        Buffer.BlockCopy(arr, start, result, 0, length);

        return result;
    }

    public static byte[] AddPssPadding(byte[] hash)
    {
        int modBits = 2048;
        int hLen = 20;
        int emLen = 256;

        int lmask = 0;
        for (int i = 0; i < 8 * emLen - (modBits - 1); i++)
        {
            lmask = (lmask >> 1) | 0x80;
        }

        // Commented out since the condition will always be false while emLen = 256 and hLen = 20
        // if (emLen < hLen + hLen + 2)
        // {
        //     return null;
        // }

        byte[] salt = new byte[hLen];
        new Random().NextBytes(salt);

        byte[] m_prime = Enumerable.Repeat((byte)0, 8).ToArray().Concat(hash).Concat(salt).ToArray();
        byte[] h = SHA1.Create().ComputeHash(m_prime);

        byte[] ps = Enumerable.Repeat((byte)0, emLen - hLen - hLen - 2).ToArray();
        byte[] db = ps.Concat(new byte[] { 0x01 }).Concat(salt).ToArray();

        byte[] dbMask = MGF1(h, emLen - hLen - 1);

        byte[] maskedDb = new byte[dbMask.Length];
        for (int i = 0; i < dbMask.Length; i++)
        {
            maskedDb[i] = (byte)(db[i] ^ dbMask[i]);
        }

        maskedDb[0] = (byte)(maskedDb[0] & ~lmask);

        byte[] padded = maskedDb.Concat(h).Concat(new byte[] { 0xBC }).ToArray();

        return padded;
    }

    public static byte[] RemoveOAEPPadding(byte[] data)
    {
        int k = 256;
        int hLen = 20;

        byte[] maskedSeed = data[1..(hLen + 1)];
        byte[] maskedDB = data[(hLen + 1)..];

        byte[] seedMask = MGF1(maskedDB, hLen);

        byte[] seed = new byte[maskedSeed.Length];
        for (int i = 0; i < maskedSeed.Length; i++)
        {
            seed[i] = (byte)(maskedSeed[i] ^ seedMask[i]);
        }

        byte[] dbMask = MGF1(seed, k - hLen - 1);

        byte[] db = new byte[maskedDB.Length];
        for (int i = 0; i < maskedDB.Length; i++)
        {
            db[i] = (byte)(maskedDB[i] ^ dbMask[i]);
        }

        int onePos = BitConverter.ToString(db[hLen..]).Replace("-", "").IndexOf("01", StringComparison.Ordinal) / 2;
        byte[] unpadded = db[(hLen + onePos + 1)..];

        return unpadded;
    }

    private static byte[] MGF1(byte[] seed, int maskLen)
    {
        SHA1 hobj = SHA1.Create();
        int hLen = hobj.HashSize / 8;
        List<byte> T = new();
        for (int i = 0; i < (int)Math.Ceiling(maskLen / (double)hLen); i++)
        {
            byte[] c = BitConverter.GetBytes(i);
            Array.Reverse(c);
            byte[] digest = hobj.ComputeHash(seed.Concat(c).ToArray());
            T.AddRange(digest);
        }

        return T.GetRange(0, maskLen).ToArray();
    }
}