napdump/ndimg/Program.cs

using System;
using System.IO;
using System.Threading.Tasks;
using System.Threading;
using System.Threading.Channels;
using System.Web;
using napdump;
using System.Runtime.Serialization.Formatters.Binary;
using System.Net;
using System.Collections.Generic;
using System.Net.Http;
using Tools.Crypto;
using System.Linq;

namespace ndimg
{
    class Program
    {
        const int MaxThreads = 3;
        static async Task<BoardInfo> ReadDump(Stream from, CancellationToken token)
        {
            var binf = new BinaryFormatter();
            await Task.Yield();
            try
            {
                var bi = (BoardInfo)binf.Deserialize(from);
                token.ThrowIfCancellationRequested();
                return bi;
            }catch(Exception ex)
            {
                throw new InvalidOperationException("Could not load dump: " + ex.Message);
            }
        }
        static readonly CancellationTokenSource cancel = new CancellationTokenSource();
        static readonly AsyncMutex semaphore = AsyncMutex.Semaphore(MaxThreads);
        static async Task DownloadImage(string url, Stream to, CancellationToken token)
        {
            using var downloader = new WebClient();
            using (await semaphore.AquireAsync(token))
            {
                await using var reader = await downloader.OpenReadTaskAsync(url);
                token.ThrowIfCancellationRequested();
                await reader.CopyToAsync(to, token);
            }
        }
        static async Task<(int Complete,int Failed)> Reader(ChannelReader<PostInfo> chan, DirectoryInfo output, CancellationToken token)
        {
            List<Task> downloaders = new List<Task>();
            int compDownloads = 0;
            await foreach(var post in chan.ReadAllAsync(token))
            {
                if (post.ImageURL != null && post.ImageURL.Length > 0)
                {
                    downloaders.Add(Task.Run(async () =>
                    {
                        var ipath = Path.Combine(output.FullName, post.PostNumber.ToString());
                        try
                        {
                            if ( DeletedKey!=null &&( post.ModLog.ImageDeleted || post.ModLog.PostDeleted))
                            {
                                await using (var tf = new TempFile())
                                {
                                    await DownloadImage(post.ImageURL, tf.Stream, token);

                                    tf.Stream.Position = 0;
                                    using (var enc = new encaes.AesEncryptor(tf.Stream))
                                    {
                                        enc.KeepAlive = true;
                                        enc.Key = DeletedKey.Value;
                                        await using (var writeStream = new FileStream(ipath, FileMode.Create))
                                        {
                                            await enc.Encrypt(writeStream, token);
                                            compDownloads += 1;
                                            Console.WriteLine($"{post.PostNumber} -> {ipath} (encrypted)");
                                            return;
                                        }
                                    }
                                }
                            }
                            else
                            {
                                await using (var writeStream = new FileStream(ipath, FileMode.Create))
                                {
                                    await DownloadImage(post.ImageURL, writeStream, token);
                                    compDownloads += 1;
                                    Console.WriteLine($"{post.PostNumber} -> {ipath}");
                                    return;
                                }
                            }
                        }
                        catch (OperationCanceledException) { }
                        catch (Exception ex)
                        {
                            Console.WriteLine($"Failed to download image for post {post.PostNumber} ({post.ImageURL}): {ex.Message}");
                        }
                        if (File.Exists(ipath)) try
                            {
                                File.Delete(ipath);
                            }
                            catch (Exception ex) { Console.WriteLine("Warning: State corrupted in file " + ipath+": " + ex.Message); }
                    }));
                }
            }
            Console.WriteLine("Waiting for downloaders...");
            await Task.WhenAll(downloaders);
            return (compDownloads, compDownloads - downloaders.Count);
        }
        public static async Task<AESKey> readKeyFromFile(string fn, CancellationToken token)
        {
            using(var fs = new FileStream(fn,FileMode.Open,FileAccess.Read))
            {
                return await encaes.AesEncryptor.LoadKey(fs, async () => await encaes.AesEncryptor.ReadPassword("Key is password protected: ", token), token);
            }
        }
        public static AESKey? DeletedKey { get; private set; } = null;
        static string isEncryptedA(ref string[] args)
        {
            for (int i = 0; i < args.Length - 1; i++)
            {
                if (args[i].ToLower() == "--deleted-key")
                {
                    var ret = args[i + 1];
                    args = args.Where((_, j) => j != i && j != (i + 1)).ToArray();
                    return ret;
                }
            }
            return null;
        }
        static async Task Main(string[] args)
        {
            if(args.Length <1)
            {
                Console.WriteLine("Usage: ndimg [--deleted-key <deleted key>] <dump> [<folder>] ");
                return;
            }

            string encFn = isEncryptedA(ref args);

            if(!File.Exists(args[0]))
            {
                Console.WriteLine("Error: dump file " + args[0] + " does not exist.");
                return;
            }


            string outputDir;
            if (args.Length < 2)
            {
                var psplit = Path.GetFileName(args[0]).Split('.');

                outputDir = Path.Join(Path.GetDirectoryName(args[0]), string.Join('.', psplit.AsMemory().Slice(0, psplit.Length - 1).ToArray()));
            }
            else
                outputDir = args[1];

            Console.WriteLine($"Downloading {args[0]} -> {outputDir}");

            DirectoryInfo output;
            if (!Directory.Exists(outputDir))
            {
                output = Directory.CreateDirectory(outputDir);
            }
            else
                output = new DirectoryInfo(outputDir);

            Console.CancelKeyPress += (o, e) =>
            {
                if (!cancel.IsCancellationRequested)
                {
                    cancel.Cancel();
                    e.Cancel = true;
                }
                else
                {
                    Console.WriteLine("Force exit");
                    Environment.Exit(-1);
                }
            };

            try
            {
                if (encFn != null)
                {
                    try
                    {
                    DeletedKey=    await readKeyFromFile(encFn, cancel.Token);
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine("Failed to read key from file " + encFn + ": " + ex.Message);
                        return;
                    }
                }
                Channel<PostInfo> posts = Channel.CreateUnbounded<PostInfo>();
                var reader = Reader(posts.Reader, output, cancel.Token);
                await using (var stream = new FileStream(args[0], FileMode.Open, FileAccess.Read))
                {
                    var board = await ReadDump(stream, cancel.Token);
                    foreach (var thread in board.Threads)
                    {
                        cancel.Token.ThrowIfCancellationRequested();

                        if (thread.IsEncrypted)
                        {
                            if (DeletedKey == null)
                            {
                                Console.WriteLine("Thread is encrypted, skipping.");
                                continue;
                            }
                            else
                            {
                                try
                                {
                                    await thread.DecryptPostAsync(DeletedKey.Value, cancel.Token);
                                }
                                catch (Exception ex)
                                {
                                    Console.WriteLine("Failed to decrypt thread, skipping: " + ex.Message);
                                    continue;
                                }
                            }
                        }

                        if(thread.IsImageEncrypted)
                        {
                            if (DeletedKey == null)
                            {
                                Console.WriteLine("Thread image is encrypted, skipping.");
                                continue;
                            }
                            else
                            {
                                try
                                {
                                    await thread.DecryptImageAsync(DeletedKey.Value, cancel.Token);
                                }
                                catch (Exception ex)
                                {
                                    Console.WriteLine("Failed to decrypt thread image, skipping: " + ex.Message);
                                    continue;
                                }
                            }
                        }

                        await posts.Writer.WriteAsync(thread, cancel.Token);
                        foreach (var post in thread.Children)
                        {
                            if(post.IsEncrypted)
                            {

                                if (DeletedKey == null)
                                {
                                    Console.WriteLine("Post is encrypted, skipping.");
                                    continue;
                                }
                                else
                                {
                                    try
                                    {
                                        await post.DecryptPostAsync(DeletedKey.Value, cancel.Token);
                                    }
                                    catch (Exception ex)
                                    {
                                        Console.WriteLine("Failed to decrypt post, skipping: " + ex.Message);
                                        continue;
                                    }
                                }
                            }
                            if (post.IsImageEncrypted)
                            {
                                if (DeletedKey == null)
                                {
                                    Console.WriteLine("Post image is encrypted, skipping.");
                                    continue;
                                }
                                else
                                {
                                    try
                                    {
                                        await post.DecryptImageAsync(DeletedKey.Value, cancel.Token);
                                    }
                                    catch (Exception ex)
                                    {
                                        Console.WriteLine("Failed to decrypt post image, skipping: " + ex.Message);
                                        continue;
                                    }
                                }
                            }

                            await posts.Writer.WriteAsync(post, cancel.Token);

                        }
                    }
                    posts.Writer.Complete();
                }
                var comp = await reader;
                Console.WriteLine($"Complete. ({comp.Complete} downloaded, {comp.Failed} failed)");
            }
            catch (Exception ex)
            {
                if (!cancel.IsCancellationRequested)
                    cancel.Cancel();
                Console.WriteLine("Error downloading: " + ex.Message);
            }
            finally
            {
                semaphore.Dispose();
                cancel.Dispose();
            }
        }
    }
}