3 files changed, 73 insertions, 7 deletions

diff --git a/html/index.html b/html/index.html
index 944005c..89882b9 100644
--- a/html/index.html
+++ b/html/index.html
@@ -61,6 +61,22 @@
       </div>
       <div id="users">
       </div>
+      <div class="modal" id="power_request_modal" tabindex="-1" aria-labelledby="power_request_modal_header" aria-hidden="true">
+        <div class="modal-dialog">
+          <div class="modal-content">
+            <div class="modal-header">
+              <h5 class="modal-title" id="power_request_modal_header">Forbidden</h5>
+              <button type="button" class="close" data-dismiss="modal" aria-label="Close">
+                <span aria-hidden="true">&times;</span>
+              </button>
+            </div>
+            <div class="modal-body">Sorry, power management is turned off on this instance.</div>
+            <div class="modal-footer">
+              <button type="button" class="btn btn-primary" data-dismiss="modal">Close</button>
+            </div>
+          </div>
+        </div>
+      </div>
     </div>
     <footer class="mt-auto py-3 border-top bg-light text-center text-muted small">
       <div class="container">
@@ -110,12 +126,29 @@ function get(url, success_callback) {
     $.get(url, success_callback).fail(dump_fail);
 }
 
+function on_power_request_forbidden() {
+    $('#power_request_modal').modal();
+}
+
+function power_request(url) {
+    $.get(url).fail(function(data) {
+        switch (data.status) {
+            case 403:
+                on_power_request_forbidden();
+                break;
+            default:
+                dump_fail(data);
+                break;
+        }
+    });
+}
+
 function reboot() {
-    get('reboot');
+    power_request('reboot');
 }
 
 function shutdown() {
-    get('poweroff');
+    power_request('poweroff');
 }
 
 function set_hostname(data) {
diff --git a/src/app.py b/src/app.py
index b83666e..7533df7 100755
--- a/src/app.py
+++ b/src/app.py
@@ -40,12 +40,10 @@
 # no more.
 
 import abc
-import cgi
 from collections import namedtuple
 from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
 from http import HTTPStatus
-import http.server
 import json
 import os
 import pwd
@@ -53,6 +51,8 @@ import shlex
 import socket
 import subprocess
 from subprocess import DEVNULL, PIPE, STDOUT
+import traceback
+import urllib.parse
 
 
 def split_by(xs, sep):
@@ -510,6 +510,17 @@ def systemd_users():
     return show_users(list_users())
 
 
+def cgi_one_value(params, name, default=None):
+    values = params.get(name, [])
+    if not values:
+        if default is None:
+            raise ValueError(f'must have at least one value: {name}')
+        return default
+    if len(values) > 1:
+        raise ValueError(f'multiple values are not supported: {name}')
+    return values[0]
+
+
 class Request(Enum):
     STATUS = 'status'
     TOP = 'top'
@@ -525,22 +536,38 @@ class Request(Enum):
             raise ValueError('HTTP path must start with a forward slash /')
         return Request(path[1:])
 
+    @staticmethod
+    def from_query_string(qs):
+        params = urllib.parse.parse_qs(qs, strict_parsing=True)
+        request = cgi_one_value(params, 'what')
+        request = Request(request)
+        request.disable_power = int(cgi_one_value(params, 'disable_power', '0'))
+        return request
+
     def process(self):
         if self is Request.STATUS:
             return Status().complete()
         if self is Request.TOP:
             return Top().complete()
+
+        if self in [Request.REBOOT, Request.POWEROFF] and self.disable_power:
+            return Response(None, HTTPStatus.FORBIDDEN)
         if self is Request.REBOOT:
             return Reboot().complete()
         if self is Request.POWEROFF:
             return Poweroff().complete()
+
         raise NotImplementedError(f'unknown request: {self}')
 
 
 def process_cgi_request():
-    params = cgi.FieldStorage()
-    what = params['what'].value
-    Request(what).process().write_as_cgi_script()
+    try:
+        request = Request.from_query_string(os.environ['QUERY_STRING'])
+        request.process().write_as_cgi_script()
+    except:
+        status = HTTPStatus.INTERNAL_SERVER_ERROR
+        response = Response(traceback.format_exc(), status)
+        response.write_as_cgi_script()
 
 
 def main():
diff --git a/src/server.py b/src/server.py
index 9ffe6f1..a0fa7c4 100755
--- a/src/server.py
+++ b/src/server.py
@@ -30,9 +30,12 @@ def default_html_dir():
 
 
 class RequestHandler(http.server.SimpleHTTPRequestHandler):
+    ARGS = None
+
     def do_GET(self):
         try:
             request = Request.from_http_path(self.path)
+            request.disable_power = RequestHandler.ARGS.disable_power
         except ValueError:
             return super().do_GET()
         try:
@@ -60,6 +63,8 @@ def parse_args(args=None):
     parser.add_argument('-p', '--port', metavar='PORT',
                         type=int, default=DEFAULT_PORT,
                         help='set port number')
+    parser.add_argument('-n', '--disable-power', action='store_true',
+                        help='disable reboot/poweroff requests')
     parser.add_argument('-d', '--dir', metavar='DIR',
                         default=default_html_dir(),
                         help='HTML directory path')
@@ -68,6 +73,7 @@ def parse_args(args=None):
 
 def main(args=None):
     args = parse_args(args)
+    RequestHandler.ARGS = args
 
     # It's a failsafe; the script is not allowed to serve a random current
     # working directory.