From 401c2d232f2137f6edaad62ba9d8950fea3da82d Mon Sep 17 00:00:00 2001
From: Timothy Baldridge <tbaldridge@gmail.com>
Date: Thu, 18 Feb 2021 17:09:12 -0700
Subject: [PATCH] Improve the GoogleDrive Downloader a bit

---
 .../Downloaders/GoogleDriveDownloader.cs        | 17 +++++++++++------
 Wabbajack.Test/DownloaderTests.cs               |  3 +++
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/Wabbajack.Lib/Downloaders/GoogleDriveDownloader.cs b/Wabbajack.Lib/Downloaders/GoogleDriveDownloader.cs
index 91b31129..902851c1 100644
--- a/Wabbajack.Lib/Downloaders/GoogleDriveDownloader.cs
+++ b/Wabbajack.Lib/Downloaders/GoogleDriveDownloader.cs
@@ -1,4 +1,5 @@
 using System.Linq;
+using System.Net.Http;
 using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
@@ -6,6 +7,7 @@ using Newtonsoft.Json;
 using Wabbajack.Common;
 using Wabbajack.Common.Exceptions;
 using Wabbajack.Common.Serialization.Json;
+using Wabbajack.Lib.Http;
 using Wabbajack.Lib.Validation;
 
 namespace Wabbajack.Lib.Downloaders
@@ -67,12 +69,15 @@ namespace Wabbajack.Lib.Downloaders
                 using var response = await client.GetAsync(initialURL);
                 if (!response.IsSuccessStatusCode)
                     throw new HttpException((int)response.StatusCode, response.ReasonPhrase ?? "Unknown");
-                var regex = new Regex("(?<=/uc\\?export=download&amp;confirm=).*(?=;id=)");
-                using var content = response.Content;
-                var confirm = regex.Match(await content.ReadAsStringAsync());
-                if (!confirm.Success)
-                    return null;
-                var url = $"https://drive.google.com/uc?export=download&confirm={confirm}&id={Id}";
+                var cookies = response.GetSetCookies();
+                var warning = cookies.FirstOrDefault(c => c.Key.StartsWith("download_warning_"));
+                response.Dispose();
+                if (warning == default)
+                {
+                    return new HTTPDownloader.State(initialURL) { Client = client };
+                }
+
+                var url = $"https://drive.google.com/uc?export=download&confirm={warning.Value}&id={Id}";
                 var httpState = new HTTPDownloader.State(url) { Client = client };
                 return httpState;
             }
diff --git a/Wabbajack.Test/DownloaderTests.cs b/Wabbajack.Test/DownloaderTests.cs
index c1a291da..95077405 100644
--- a/Wabbajack.Test/DownloaderTests.cs
+++ b/Wabbajack.Test/DownloaderTests.cs
@@ -113,6 +113,9 @@ namespace Wabbajack.Test
             Assert.Equal(Hash.FromBase64("eSIyd+KOG3s="), await filename.Path.FileHashAsync());
 
             Assert.Equal("Cheese for Everyone!", await filename.Path.ReadAllTextAsync());
+
+            var newState = (AbstractDownloadState)new GoogleDriveDownloader.State("1Q_CdeYJStfoTZFLZ79RRVkxI2c_cG0dg");
+            Assert.True(await newState.Verify(new Archive(newState) {Size = 0}));
         }
 
         [Fact]