From 941f3f0934cbe6ec40b7da05f01242489caf8a83 Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Tue, 1 Aug 2017 15:12:12 +0900
Subject: [PATCH] Tidy up osu-stable import process

Now can locate any osu-stable installation using registry lookup (with ample fallbacks). Also uses a much more controlled access method via StableStorage.
---
 osu.Desktop/OsuGameDesktop.cs       | 52 +++++++++++++++++++++++++++++
 osu.Game/Beatmaps/BeatmapManager.cs | 18 ++++++----
 osu.Game/OsuGame.cs                 |  4 +++
 3 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/osu.Desktop/OsuGameDesktop.cs b/osu.Desktop/OsuGameDesktop.cs
index bd5c6c6790..f1427d2861 100644
--- a/osu.Desktop/OsuGameDesktop.cs
+++ b/osu.Desktop/OsuGameDesktop.cs
@@ -1,6 +1,7 @@
 // Copyright (c) 2007-2017 ppy Pty Ltd <contact@ppy.sh>.
 // Licensed under the MIT Licence - https://raw.githubusercontent.com/ppy/osu/master/LICENCE
 
+using System;
 using osu.Game;
 using System.Linq;
 using System.Windows.Forms;
@@ -11,6 +12,7 @@ using System.Reflection;
 using System.Drawing;
 using System.IO;
 using System.Threading.Tasks;
+using Microsoft.Win32;
 using osu.Framework.Graphics.Containers;
 using osu.Game.Screens.Menu;
 
@@ -30,6 +32,56 @@ namespace osu.Desktop
             };
         }
 
+        public override Storage GetStorageForStableInstall()
+        {
+            try
+            {
+                return new StableStorage();
+            }
+            catch
+            {
+                return null;
+            }
+        }
+
+        /// <summary>
+        /// A method of accessing an osu-stable install in a controlled fashion.
+        /// </summary>
+        private class StableStorage : DesktopStorage
+        {
+            protected override string LocateBasePath()
+            {
+                string stableInstallPath;
+
+                try
+                {
+                    using (RegistryKey key = Registry.ClassesRoot.OpenSubKey("osu"))
+                        stableInstallPath = key?.OpenSubKey(@"shell\open\command")?.GetValue(String.Empty).ToString().Split('"')[1].Replace("osu!.exe", "");
+
+                    if (Directory.Exists(stableInstallPath))
+                        return stableInstallPath;
+                }
+                catch
+                {
+                }
+
+                stableInstallPath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData), @"osu!", "Songs");
+                if (Directory.Exists(stableInstallPath))
+                    return stableInstallPath;
+
+                stableInstallPath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".osu", "Songs");
+                if (Directory.Exists(stableInstallPath))
+                    return stableInstallPath;
+
+                return null;
+            }
+
+            public StableStorage()
+                : base(string.Empty)
+            {
+            }
+        }
+
         protected override void LoadComplete()
         {
             base.LoadComplete();
diff --git a/osu.Game/Beatmaps/BeatmapManager.cs b/osu.Game/Beatmaps/BeatmapManager.cs
index 5cc2a2a430..6059db0a36 100644
--- a/osu.Game/Beatmaps/BeatmapManager.cs
+++ b/osu.Game/Beatmaps/BeatmapManager.cs
@@ -59,6 +59,11 @@ namespace osu.Game.Beatmaps
         /// </summary>
         public Action<Notification> PostNotification { private get; set; }
 
+        /// <summary>
+        /// Set a storage with access to an osu-stable install for import purposes.
+        /// </summary>
+        public Func<Storage> GetStableStorage { private get; set; }
+
         public BeatmapManager(Storage storage, FileStore files, SQLiteConnection connection, RulesetStore rulesets, IIpcHost importHost = null)
         {
             beatmaps = new BeatmapStore(connection);
@@ -451,19 +456,20 @@ namespace osu.Game.Beatmaps
             }
         }
 
+        /// <summary>
+        /// This is a temporary method and will likely be replaced by a full-fledged (and more correctly placed) migration process in the future.
+        /// </summary>
         public void ImportFromStable()
         {
-            string stableInstallPath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData), @"osu!", "Songs");
-            if (!Directory.Exists(stableInstallPath))
-                stableInstallPath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".osu", "Songs");
+            var stable = GetStableStorage?.Invoke();
 
-            if (!Directory.Exists(stableInstallPath))
+            if (stable == null)
             {
-                Logger.Log("Couldn't find an osu!stable installation!", LoggingTarget.Information, LogLevel.Error);
+                Logger.Log("No osu!stable installation available!", LoggingTarget.Information, LogLevel.Error);
                 return;
             }
 
-            Import(Directory.GetDirectories(stableInstallPath));
+            Import(stable.GetDirectories("Songs"));
         }
 
         public void DeleteAll()
diff --git a/osu.Game/OsuGame.cs b/osu.Game/OsuGame.cs
index a70a7b7d97..8d8c5cf26e 100644
--- a/osu.Game/OsuGame.cs
+++ b/osu.Game/OsuGame.cs
@@ -20,6 +20,7 @@ using osu.Game.Screens.Menu;
 using OpenTK;
 using System.Linq;
 using System.Threading.Tasks;
+using osu.Framework.Platform;
 using osu.Framework.Threading;
 using osu.Game.Graphics;
 using osu.Game.Rulesets.Scoring;
@@ -47,6 +48,8 @@ namespace osu.Game
 
         private UserProfileOverlay userProfile;
 
+        public virtual Storage GetStorageForStableInstall() => null;
+
         private Intro intro
         {
             get
@@ -151,6 +154,7 @@ namespace osu.Game
 
             // hook up notifications to components.
             BeatmapManager.PostNotification = n => notificationOverlay?.Post(n);
+            BeatmapManager.GetStableStorage = GetStorageForStableInstall;
 
             AddRange(new Drawable[] {
                 new VolumeControlReceptor