From 0bdb48bcac9244df6700410b9cf4ec581dbbda38 Mon Sep 17 00:00:00 2001
From: Laszlo Gazdag <lazlyhu@users.noreply.github.com>
Date: Wed, 10 Nov 2021 20:31:22 +0100
Subject: [PATCH] Make OTA function switchable in web_server component (#2685)

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
---
 esphome/components/web_server/__init__.py    |  3 +++
 esphome/components/web_server/web_server.cpp | 16 +++++++++++-----
 esphome/components/web_server/web_server.h   |  7 +++++++
 tests/test1.yaml                             |  1 +
 tests/test4.yaml                             |  1 +
 5 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/esphome/components/web_server/__init__.py b/esphome/components/web_server/__init__.py
index ba2d866593..61b1fa5ad6 100644
--- a/esphome/components/web_server/__init__.py
+++ b/esphome/components/web_server/__init__.py
@@ -12,6 +12,7 @@ from esphome.const import (
     CONF_AUTH,
     CONF_USERNAME,
     CONF_PASSWORD,
+    CONF_OTA,
 )
 from esphome.core import CORE, coroutine_with_priority
 
@@ -41,6 +42,7 @@ CONFIG_SCHEMA = cv.Schema(
         cv.GenerateID(CONF_WEB_SERVER_BASE_ID): cv.use_id(
             web_server_base.WebServerBase
         ),
+        cv.Optional(CONF_OTA, default=True): cv.boolean,
     }
 ).extend(cv.COMPONENT_SCHEMA)
 
@@ -57,6 +59,7 @@ async def to_code(config):
     cg.add_define("USE_WEBSERVER")
     cg.add(var.set_css_url(config[CONF_CSS_URL]))
     cg.add(var.set_js_url(config[CONF_JS_URL]))
+    cg.add(var.set_allow_ota(config[CONF_OTA]))
     if CONF_AUTH in config:
         cg.add(paren.set_auth_username(config[CONF_AUTH][CONF_USERNAME]))
         cg.add(paren.set_auth_password(config[CONF_AUTH][CONF_PASSWORD]))
diff --git a/esphome/components/web_server/web_server.cpp b/esphome/components/web_server/web_server.cpp
index 17b17fcc3c..6f47f460af 100644
--- a/esphome/components/web_server/web_server.cpp
+++ b/esphome/components/web_server/web_server.cpp
@@ -152,7 +152,9 @@ void WebServer::setup() {
 #endif
   this->base_->add_handler(&this->events_);
   this->base_->add_handler(this);
-  this->base_->add_ota_handler();
+
+  if (this->allow_ota_)
+    this->base_->add_ota_handler();
 
   this->set_interval(10000, [this]() { this->events_.send("", "ping", millis(), 30000); });
 }
@@ -240,10 +242,14 @@ void WebServer::handle_index_request(AsyncWebServerRequest *request) {
 #endif
 
   stream->print(F("</tbody></table><p>See <a href=\"https://esphome.io/web-api/index.html\">ESPHome Web API</a> for "
-                  "REST API documentation.</p>"
-                  "<h2>OTA Update</h2><form method=\"POST\" action=\"/update\" enctype=\"multipart/form-data\"><input "
-                  "type=\"file\" name=\"update\"><input type=\"submit\" value=\"Update\"></form>"
-                  "<h2>Debug Log</h2><pre id=\"log\"></pre>"));
+                  "REST API documentation.</p>"));
+  if (this->allow_ota_) {
+    stream->print(
+        F("<h2>OTA Update</h2><form method=\"POST\" action=\"/update\" enctype=\"multipart/form-data\"><input "
+          "type=\"file\" name=\"update\"><input type=\"submit\" value=\"Update\"></form>"));
+  }
+  stream->print(F("<h2>Debug Log</h2><pre id=\"log\"></pre>"));
+
 #ifdef WEBSERVER_JS_INCLUDE
   if (this->js_include_ != nullptr) {
     stream->print(F("<script src=\"/0.js\"></script>"));
diff --git a/esphome/components/web_server/web_server.h b/esphome/components/web_server/web_server.h
index 021d5a0646..cdfec51cf1 100644
--- a/esphome/components/web_server/web_server.h
+++ b/esphome/components/web_server/web_server.h
@@ -58,6 +58,12 @@ class WebServer : public Controller, public Component, public AsyncWebHandler {
    */
   void set_js_include(const char *js_include);
 
+  /** Set whether or not the webserver should expose the OTA form and handler.
+   *
+   * @param allow_ota.
+   */
+  void set_allow_ota(bool allow_ota) { this->allow_ota_ = allow_ota; }
+
   // ========== INTERNAL METHODS ==========
   // (In most use cases you won't need these)
   /// Setup the internal web server and register handlers.
@@ -182,6 +188,7 @@ class WebServer : public Controller, public Component, public AsyncWebHandler {
   const char *css_include_{nullptr};
   const char *js_url_{nullptr};
   const char *js_include_{nullptr};
+  bool allow_ota_{true};
 };
 
 }  // namespace web_server
diff --git a/tests/test1.yaml b/tests/test1.yaml
index 585e01635f..dee7493bf2 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -234,6 +234,7 @@ logger:
 
 web_server:
   port: 8080
+  ota: true
   css_url: https://esphome.io/_static/webserver-v1.min.css
   js_url: https://esphome.io/_static/webserver-v1.min.js
 
diff --git a/tests/test4.yaml b/tests/test4.yaml
index 4228c7494c..938145235a 100644
--- a/tests/test4.yaml
+++ b/tests/test4.yaml
@@ -45,6 +45,7 @@ logger:
   level: DEBUG
 
 web_server:
+  ota: false
   auth:
     username: admin
     password: admin