4 files changed, 44 insertions, 34 deletions

diff --git a/gui/action.cpp b/gui/action.cpp
index 8ea1dae82..8315ab2e9 100644
--- a/gui/action.cpp
+++ b/gui/action.cpp
@@ -1354,44 +1354,18 @@ int GUIAction::decrypt(std::string arg)
 		if (op_status != 0)
 			op_status = 1;
 		else {
-			int load_theme = 1;
 
 			DataManager::SetValue(TW_IS_ENCRYPTED, 0);
 
-			if (load_theme) {
-				int has_datamedia;
+			int has_datamedia;
 
-				// Check for a custom theme and load it if exists
-				DataManager::GetValue(TW_HAS_DATA_MEDIA, has_datamedia);
-				if (has_datamedia != 0) {
-					struct stat st;
-					int check = 0;
-					std::string theme_path;
-
-					if (tw_get_default_metadata(DataManager::GetSettingsStoragePath().c_str()) != 0) {
-						LOGERR("Failed to get default contexts and file mode for storage files.\n");
-					} else {
-						LOGINFO("Got default contexts and file mode for storage files.\n");
-					}
-
-					theme_path = DataManager::GetSettingsStoragePath();
-					if (PartitionManager.Mount_By_Path(theme_path.c_str(), 1) < 0) {
-						LOGERR("Unable to mount %s during reload function startup.\n", theme_path.c_str());
-						check = 1;
-					}
-
-					theme_path += "/TWRP/theme/ui.zip";
-					if (check == 0 && stat(theme_path.c_str(), &st) == 0) {
-						if (PageManager::ReloadPackage("TWRP", theme_path) != 0)
-						{
-							// Loading the custom theme failed - try loading the stock theme
-							LOGINFO("Attempting to reload stock theme...\n");
-							if (PageManager::ReloadPackage("TWRP", "/res/ui.xml"))
-							{
-								LOGERR("Failed to load base packages.\n");
-							}
-						}
-					}
+			// Check for a custom theme and load it if exists
+			DataManager::GetValue(TW_HAS_DATA_MEDIA, has_datamedia);
+			if (has_datamedia != 0) {
+				if (tw_get_default_metadata(DataManager::GetSettingsStoragePath().c_str()) != 0) {
+					LOGERR("Failed to get default contexts and file mode for storage files.\n");
+				} else {
+					LOGINFO("Got default contexts and file mode for storage files.\n");
 				}
 			}
 		}
diff --git a/gui/gui.cpp b/gui/gui.cpp
index edf2d7671..61e5de1fb 100644
--- a/gui/gui.cpp
+++ b/gui/gui.cpp
@@ -805,6 +805,38 @@ error:
 	return -1;
 }
 
+extern "C" int gui_loadCustomResources(void)
+{
+#ifndef TW_OEM_BUILD
+	if (!PartitionManager.Mount_Settings_Storage(false)) {
+		LOGERR("Unable to mount settings storage during GUI startup.\n");
+		return -1;
+	}
+
+	std::string theme_path = DataManager::GetSettingsStoragePath();
+	theme_path += "/TWRP/theme/ui.zip";
+	// Check for a custom theme
+	if (TWFunc::Path_Exists(theme_path)) {
+		// There is a custom theme, try to load it
+		if (PageManager::ReloadPackage("TWRP", theme_path)) {
+			// Custom theme failed to load, try to load stock theme
+			if (PageManager::ReloadPackage("TWRP", "/res/ui.xml")) {
+				LOGERR("Failed to load base packages.\n");
+				goto error;
+			}
+		}
+	}
+	// Set the default package
+	PageManager::SelectPackage("TWRP");
+#endif
+	return 0;
+
+error:
+	LOGERR("An internal error has occurred: unable to load theme.\n");
+	gGuiInitialized = 0;
+	return -1;
+}
+
 extern "C" int gui_start(void)
 {
 	if (!gGuiInitialized)
diff --git a/gui/gui.h b/gui/gui.h
index 8e0371ef2..21c460ebe 100644
--- a/gui/gui.h
+++ b/gui/gui.h
@@ -24,6 +24,7 @@
 int gui_console_only();
 int gui_init();
 int gui_loadResources();
+int gui_loadCustomResources();
 int gui_start();
 int gui_startPage(const char* page_name);
 void gui_print(const char *fmt, ...);
diff --git a/twrp.cpp b/twrp.cpp
index 35404c34d..3489f8cfc 100644
--- a/twrp.cpp
+++ b/twrp.cpp
@@ -278,6 +278,9 @@ int main(int argc, char **argv) {
 		LOGINFO("Is encrypted, do decrypt page first\n");
 		if (gui_startPage("decrypt") != 0) {
 			LOGERR("Failed to start decrypt GUI page.\n");
+		} else {
+			// Check for and load custom theme if present
+			gui_loadCustomResources();
 		}
 	} else if (datamedia) {
 		if (tw_get_default_metadata(DataManager::GetSettingsStoragePath().c_str()) != 0) {