<commit_before>﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Security;
using System.Security.Cryptography;

namespace Vault
{
    public unsafe static class Security
    {
        const int DEFAULT_ITERATIONS = 1675; //Randomly picked number
        const int DEFAULT_SALTSIZE = 8;

        public static void MergeFile(IDictionary<string, SecureString> values, string path, byte[] password, ushort saltSize = DEFAULT_SALTSIZE, int iterations = DEFAULT_ITERATIONS, bool encryptResult = true)
        {
            var dictionary = DecryptFile(path, password, iterations, encryptResult);
            var original = dictionary.Values.ToArray();

            foreach (var kvp in values)
                dictionary[kvp.Key] = kvp.Value;

            EncryptFile(dictionary, path, password, saltSize, iterations, encryptResult);

            //Clean up decrypted keys, make user clean up their own.
            foreach (var secureString in original)
                secureString.Dispose();
        }

        public static void EncryptFile(IDictionary<string, SecureString> values, string path, byte[] password, ushort saltSize = DEFAULT_SALTSIZE, int iterations = DEFAULT_ITERATIONS, bool encryptResult = true)
        {
            var bytes = EncryptDictionary(values, password, saltSize, iterations, encryptResult);
            File.WriteAllBytes(path, bytes);
            bytes.Clear();
        }

        public static IDictionary<string, SecureString> DecryptFile(string path, byte[] password, int iterations = DEFAULT_ITERATIONS, bool inputIsEncrypted = true)
        {
            if (!File.Exists(path)) return new Dictionary<string, SecureString>();

            var bytes = File.ReadAllBytes(path);
            var result = DecryptDictionary(bytes, password, iterations, inputIsEncrypted);
            bytes.Clear();

            return result;
        }

        public static byte[] EncryptDictionary(IDictionary<string, SecureString> values, byte[] password, ushort saltSize = DEFAULT_SALTSIZE, int iterations = DEFAULT_ITERATIONS, bool encryptResult = true)
        {
            if (values.Count == 0) return new byte[0];

            using (var stream = new MemoryStream())
            {
                foreach (var kvp in values)
                {
                    var key = kvp.Key;
                    var keyArray = new byte[key.Length * sizeof(char)];

                    fixed (void* keyPtr = key)
                    fixed (void* destPtr = keyArray)
                        UnsafeNativeMethods.memcpy(destPtr, keyPtr, keyArray.Length);

                    var encrypted = EncryptSecureString(kvp.Value, password, saltSize, iterations);

                    var index = new byte[sizeof(ushort)];

                    fixed (void* ptr = index)
                        *((ushort*)ptr) = (ushort)keyArray.Length;
                    stream.Write(index, 0, sizeof(ushort));

                    fixed (void* ptr = index)
                        *((ushort*)ptr) = (ushort)encrypted.Length;
                    stream.Write(index, 0, sizeof(ushort));

                    stream.Write(keyArray, 0, keyArray.Length);
                    stream.Write(encrypted, 0, encrypted.Length);

                    encrypted.Clear();
                }

                var result = stream.ToArray();
                if (!encryptResult) return result;

                var encrypt = Encrypt(result, password);
                result.Clear();
                return encrypt;
            }
        }

        public static SecureString DecryptDictionary(byte[] input, string dictionaryKey, byte[] password, int iterations = DEFAULT_ITERATIONS, bool inputIsEncrypted = true, StringComparer comparer = null)
        {<commit_msg>Writing index file for fast key lookup<commit_after>﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Security;
using System.Security.Cryptography;

namespace Vault
{
    public unsafe static class Security
    {
        const int DEFAULT_ITERATIONS = 1675; //Randomly picked number
        const int DEFAULT_SALTSIZE = 8;

        public static void MergeFile(IDictionary<string, SecureString> values, string path, byte[] password, ushort saltSize = DEFAULT_SALTSIZE, int iterations = DEFAULT_ITERATIONS, bool encryptResult = true)
        {
            var dictionary = DecryptFile(path, password, iterations, encryptResult);
            var original = dictionary.Values.ToArray();

            foreach (var kvp in values)
                dictionary[kvp.Key] = kvp.Value;

            EncryptFile(dictionary, path, password, saltSize, iterations, encryptResult);

            //Clean up decrypted keys, make user clean up their own.
            foreach (var secureString in original)
                secureString.Dispose();
        }
        
        public static void EncryptFile(IDictionary<string, SecureString> values, string path, byte[] password, bool writeOffsets = true, ushort saltSize = DEFAULT_SALTSIZE, int iterations = DEFAULT_ITERATIONS, bool encryptResult = true)
        {
            byte[] bytes;

            if (writeOffsets)
            {
                byte[] offsets;
                bytes = EncryptDictionary(values, password, out offsets, saltSize, iterations, encryptResult);



                var idx = Path.Combine(Path.GetDirectoryName(path), Path.GetFileNameWithoutExtension(path)) + ".idx";
                File.WriteAllBytes(idx, offsets);
                offsets.Clear();
            }
            else
            {
                bytes = EncryptDictionary(values, password, saltSize, iterations, encryptResult);
            }

            File.WriteAllBytes(path, bytes);

            bytes.Clear();
        }

        public static IDictionary<string, SecureString> DecryptFile(string path, byte[] password, int iterations = DEFAULT_ITERATIONS, bool inputIsEncrypted = true)
        {
            if (!File.Exists(path)) return new Dictionary<string, SecureString>();

            //TODO: Check for idx file

            var bytes = File.ReadAllBytes(path);
            var result = DecryptDictionary(bytes, password, iterations, inputIsEncrypted);
            bytes.Clear();

            return result;
        }

        public static byte[] EncryptDictionary(IDictionary<string, SecureString> values, byte[] password, ushort saltSize = DEFAULT_SALTSIZE, int iterations = DEFAULT_ITERATIONS, bool encryptResult = true)
        {
            byte[] offsets;
            var result = EncryptDictionary(values, password, out offsets, saltSize, iterations, encryptResult);

            offsets.Clear();

            return result;
        }

        public static byte[] EncryptDictionary(IDictionary<string, SecureString> values, byte[] password, out byte[] offsets, ushort saltSize = DEFAULT_SALTSIZE, int iterations = DEFAULT_ITERATIONS, bool encryptResult = true)
        {
            if (values.Count == 0) return offsets = new byte[0];

            offsets = new byte[values.Count * sizeof(ushort)];

            using (var stream = new MemoryStream())
            {
                var counter = 0;
                ushort offset = 0;

                foreach (var kvp in values)
                {
                    var key = kvp.Key;
                    var keyArray = new byte[key.Length * sizeof(char)];

                    fixed (void* keyPtr = key)
                    fixed (void* destPtr = keyArray)
                        memcpy(destPtr, keyPtr, keyArray.Length);
                    var encrypted = EncryptSecureString(kvp.Value, password, saltSize, iterations);

                    var index = new byte[sizeof(ushort)];

                    fixed (void* ptr = index)
                        *((ushort*)ptr) = (ushort)keyArray.Length;
                    stream.Write(index, 0, sizeof(ushort));

                    fixed (void* ptr = index)
                        *((ushort*)ptr) = (ushort)encrypted.Length;
                    stream.Write(index, 0, sizeof(ushort));

                    stream.Write(keyArray, 0, keyArray.Length);
                    stream.Write(encrypted, 0, encrypted.Length);

                    encrypted.Clear();

                    fixed (byte* b = &offsets[counter])
                        *((ushort*)b) = offset;

                    counter += sizeof(ushort);

                    offset = (ushort)(offset + (sizeof(ushort) * 2) + keyArray.Length + encrypted.Length);
                }

                var result = stream.ToArray();

                if (!encryptResult) return result;

                var originalOffsets = offsets;
                offsets = Encrypt(offsets, password, saltSize, iterations);
                originalOffsets.Clear();

                var encrypt = Encrypt(result, password);
                result.Clear();

                return encrypt;
            }
        }

        public static SecureString DecryptDictionary(byte[] input, string dictionaryKey, byte[] password, int iterations = DEFAULT_ITERATIONS, bool inputIsEncrypted = true, StringComparer comparer = null)
        {
{'fraction_non_alphanumeric': 0.08910785048730856, 'fraction_numerical': 0.002677519545892685, 'mean_word_length': 2.0466557911908647, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4355384', 'n_tokens_mistral': 2450, 'n_tokens_neox': 2200, 'n_words': 650}