From 4a2d8c22cad1886b54df539306f45590fb34cf07 Mon Sep 17 00:00:00 2001 From: Thomas Cheyney Date: Wed, 24 Nov 2021 11:06:22 +0000 Subject: [PATCH] Add parameter to allow downloads to prefer to use the name provided by the server for downloaded files --- .../Classes/External Interfaces/Playlists.cs | 8 ++--- .../Classes/External Interfaces/Utils.cs | 33 +++++++++---------- ModAssistant/Classes/Updater.cs | 2 +- ModAssistant/Classes/Utils.cs | 19 +++++++++-- 4 files changed, 36 insertions(+), 26 deletions(-) diff --git a/ModAssistant/Classes/External Interfaces/Playlists.cs b/ModAssistant/Classes/External Interfaces/Playlists.cs index 93f3d5a0..6442c700 100644 --- a/ModAssistant/Classes/External Interfaces/Playlists.cs +++ b/ModAssistant/Classes/External Interfaces/Playlists.cs @@ -2,7 +2,6 @@ using System.IO; using System.Linq; using System.Threading.Tasks; -using System.Web; using System.Windows; using static ModAssistant.Http; @@ -34,14 +33,11 @@ public static async Task DownloadAll(Uri uri) public static async Task Get(Uri url) { - string filename = HttpUtility.UrlDecode(url.Segments.Last()); - string absolutePath = Path.Combine(BeatSaberPath, PlaylistsFolder, filename); try { CreatePlaylistsFolder(); - await Utils.DownloadAsset(url.ToString(), PlaylistsFolder, filename); - - return absolutePath; + string filename = await Utils.DownloadAsset(url.ToString(), PlaylistsFolder, preferContentDisposition: true); + return Path.Combine(BeatSaberPath, PlaylistsFolder, filename); } catch { diff --git a/ModAssistant/Classes/External Interfaces/Utils.cs b/ModAssistant/Classes/External Interfaces/Utils.cs index e9a0cbf5..c7a1b59d 100644 --- a/ModAssistant/Classes/External Interfaces/Utils.cs +++ b/ModAssistant/Classes/External Interfaces/Utils.cs @@ -29,17 +29,7 @@ public static void SetMessage(string message) } } - public static async Task DownloadAsset(string link, string folder, bool showNotifcation, string fileName = null) - { - await DownloadAsset(link, folder, fileName, null, showNotifcation); - } - - public static async Task DownloadAsset(string link, string folder, string fileName = null, string displayName = null) - { - await DownloadAsset(link, folder, fileName, displayName, true); - } - - public static async Task DownloadAsset(string link, string folder, string fileName, string displayName, bool showNotification, bool beatsaver = false) + public static async Task DownloadAsset(string link, string folder, string fileName = null, string displayName = null, bool showNotification = true, bool beatsaver = false, bool preferContentDisposition = false) { if (string.IsNullOrEmpty(BeatSaberPath)) { @@ -47,32 +37,41 @@ public static async Task DownloadAsset(string link, string folder, string fileNa } try { - Directory.CreateDirectory(Path.Combine(BeatSaberPath, folder)); + var parentDir = Path.Combine(BeatSaberPath, folder); + Directory.CreateDirectory(parentDir); + if (string.IsNullOrEmpty(fileName)) { - fileName = WebUtility.UrlDecode(Path.Combine(BeatSaberPath, folder, new Uri(link).Segments.Last())); + fileName = new Uri(link).Segments.Last(); + } + + if (beatsaver) + { + fileName = WebUtility.UrlDecode(Path.Combine(parentDir, fileName)); + await BeatSaver.Download(link, fileName); } else { - fileName = WebUtility.UrlDecode(Path.Combine(BeatSaberPath, folder, fileName)); + fileName = await ModAssistant.Utils.Download(link, parentDir, fileName, preferContentDisposition); } + if (string.IsNullOrEmpty(displayName)) { displayName = Path.GetFileNameWithoutExtension(fileName); } - if (beatsaver) await BeatSaver.Download(link, fileName); - else await ModAssistant.Utils.Download(link, fileName); - if (showNotification) { SetMessage(string.Format((string)Application.Current.FindResource("OneClick:InstalledAsset"), displayName)); } + + return fileName; } catch { SetMessage((string)Application.Current.FindResource("OneClick:AssetInstallFailed")); App.CloseWindowOnFinish = false; + return null; } } } diff --git a/ModAssistant/Classes/Updater.cs b/ModAssistant/Classes/Updater.cs index 7f88eb1d..4282826a 100644 --- a/ModAssistant/Classes/Updater.cs +++ b/ModAssistant/Classes/Updater.cs @@ -77,7 +77,7 @@ public static async Task StartUpdate() File.Move(Utils.ExePath, OldExe); - await Utils.Download(DownloadLink, NewExe); + await Utils.Download(DownloadLink, "", NewExe); RunNew(); } } diff --git a/ModAssistant/Classes/Utils.cs b/ModAssistant/Classes/Utils.cs index 886e6700..9711a538 100644 --- a/ModAssistant/Classes/Utils.cs +++ b/ModAssistant/Classes/Utils.cs @@ -5,6 +5,7 @@ using System.IO; using System.Linq; using System.Management; +using System.Net; using System.Security.Cryptography; using System.Security.Principal; using System.Text; @@ -438,14 +439,28 @@ public static void Log(string message, string severity = "LOG") File.AppendAllText(logFile, $"[{DateTime.UtcNow:yyyy-mm-dd HH:mm:ss.ffffff}][{severity.ToUpper()}] {message}\n"); } - public static async Task Download(string link, string output) + public static async Task Download(string link, string folder, string output, bool preferContentDisposition = false) { var resp = await HttpClient.GetAsync(link); + var cdFilename = resp.Content.Headers.ContentDisposition.FileName.Trim('"'); + // Prevent path traversal + if (cdFilename.IndexOfAny(Path.GetInvalidFileNameChars()) >= 0) + { + cdFilename = null; + } + + var filename = WebUtility.UrlDecode(Path.Combine( + folder, + (preferContentDisposition ? cdFilename : null) ?? output + )); + using (var stream = await resp.Content.ReadAsStreamAsync()) - using (var fs = new FileStream(output, FileMode.OpenOrCreate, FileAccess.Write)) + using (var fs = new FileStream(filename, FileMode.OpenOrCreate, FileAccess.Write)) { await stream.CopyToAsync(fs); } + + return filename; } private delegate void ShowMessageBoxDelegate(string Message, string Caption);