1 files changed, 23 insertions, 6 deletions

diff --git a/modules/gallery/controllers/upgrader.php b/modules/gallery/controllers/upgrader.php
index 0d5bb4f6..5eb96fdd 100644
--- a/modules/gallery/controllers/upgrader.php
+++ b/modules/gallery/controllers/upgrader.php
@@ -19,20 +19,33 @@
  */
 class Upgrader_Controller extends Controller {
   public function index() {
-    // Todo: give the admin a chance to log in here
-    if (!user::active()->admin) {
-      access::forbidden();
+    $session = Session::instance();
+
+    // Make sure we have an upgrade token
+    if (!($upgrade_token = $session->get("upgrade_token", null))) {
+      $session->set("upgrade_token", $upgrade_token = md5(rand()));
+    }
+
+    // If the upgrade token exists, then bless this session
+    if (file_exists(TMPPATH . $upgrade_token)) {
+      $session->set("can_upgrade", true);
+      @unlink(TMPPATH . $upgrade_token);
     }
 
     $view = new View("upgrader.html");
+    $view->can_upgrade = user::active()->admin || $session->get("can_upgrade");
+    $view->upgrade_token = $upgrade_token;
     $view->available = module::available();
     $view->done = Input::instance()->get("done");
     print $view;
   }
 
   public function upgrade() {
-    // Todo: give the admin a chance to log in here
-    if (!user::active()->admin) {
+    if (php_sapi_name() == "cli") {
+      // @todo this may screw up some module installers, but we don't have a better answer at
+      // this time.
+      $_SERVER["HTTP_HOST"] = "example.com";
+    } else if (!user::active()->admin && !Session::instance()->get("can_upgrade", false)) {
       access::forbidden();
     }
 
@@ -51,6 +64,10 @@ class Upgrader_Controller extends Controller {
       }
     }
 
-    url::redirect("upgrader?done=1");
+    if (php_sapi_name() == "cli") {
+      print "Upgrade complete\n";
+    } else {
+      url::redirect("upgrader?done=1");
+    }
   }
 }