OTA firmware MD5 check + password support for esp-idf (#2507)

Co-authored-by: Maurice Makaay <account-github@makaay.nl>

esphome/components/ota/__init__.py

@@ -15,7 +15,7 @@ from esphome.core import CORE, coroutine_with_priority
 
 CODEOWNERS = ["@esphome/core"]
 DEPENDENCIES = ["network"]
-AUTO_LOAD = ["socket"]
+AUTO_LOAD = ["socket", "md5"]
 
 CONF_ON_STATE_CHANGE = "on_state_change"
 CONF_ON_BEGIN = "on_begin"
@@ -35,20 +35,12 @@ OTAEndTrigger = ota_ns.class_("OTAEndTrigger", automation.Trigger.template())
 OTAErrorTrigger = ota_ns.class_("OTAErrorTrigger", automation.Trigger.template())
 
 
-def validate_password_support(value):
-    if CORE.using_arduino:
-        return value
-    if CORE.using_esp_idf:
-        raise cv.Invalid("Password support is not implemented yet for ESP-IDF")
-    raise NotImplementedError
-
-
 CONFIG_SCHEMA = cv.Schema(
     {
         cv.GenerateID(): cv.declare_id(OTAComponent),
         cv.Optional(CONF_SAFE_MODE, default=True): cv.boolean,
         cv.SplitDefault(CONF_PORT, esp8266=8266, esp32=3232): cv.port,
-        cv.Optional(CONF_PASSWORD): cv.All(cv.string, validate_password_support),
+        cv.Optional(CONF_PASSWORD): cv.string,
         cv.Optional(
             CONF_REBOOT_TIMEOUT, default="5min"
         ): cv.positive_time_period_milliseconds,

esphome/components/ota/ota_backend_arduino_esp32.h

@@ -9,6 +9,7 @@ namespace esphome {
 namespace ota {
 
 class ArduinoESP32OTABackend : public OTABackend {
+ public:
   OTAResponseTypes begin(size_t image_size) override;
   void set_update_md5(const char *md5) override;
   OTAResponseTypes write(uint8_t *data, size_t len) override;

esphome/components/ota/ota_backend_esp_idf.cpp

@@ -4,6 +4,7 @@
 #include "ota_backend_esp_idf.h"
 #include "ota_component.h"
 #include <esp_ota_ops.h>
+#include "esphome/components/md5/md5.h"
 
 namespace esphome {
 namespace ota {
@@ -24,15 +25,15 @@ OTAResponseTypes IDFOTABackend::begin(size_t image_size) {
     }
     return OTA_RESPONSE_ERROR_UNKNOWN;
   }
+  this->md5_.init();
   return OTA_RESPONSE_OK;
 }
 
-void IDFOTABackend::set_update_md5(const char *md5) {
-  // pass
-}
+void IDFOTABackend::set_update_md5(const char *expected_md5) { memcpy(this->expected_bin_md5_, expected_md5, 32); }
 
 OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
   esp_err_t err = esp_ota_write(this->update_handle_, data, len);
+  this->md5_.add(data, len);
   if (err != ESP_OK) {
     if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
       return OTA_RESPONSE_ERROR_MAGIC;
@@ -45,6 +46,11 @@ OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
 }
 
 OTAResponseTypes IDFOTABackend::end() {
+  this->md5_.calculate();
+  if (!this->md5_.equals_hex(this->expected_bin_md5_)) {
+    this->abort();
+    return OTA_RESPONSE_ERROR_UPDATE_END;
+  }
   esp_err_t err = esp_ota_end(this->update_handle_);
   this->update_handle_ = 0;
   if (err == ESP_OK) {

esphome/components/ota/ota_backend_esp_idf.h

@@ -5,6 +5,7 @@
 #include "ota_component.h"
 #include "ota_backend.h"
 #include <esp_ota_ops.h>
+#include "esphome/components/md5/md5.h"
 
 namespace esphome {
 namespace ota {
@@ -20,6 +21,8 @@ class IDFOTABackend : public OTABackend {
  private:
   esp_ota_handle_t update_handle_{0};
   const esp_partition_t *partition_;
+  md5::MD5Digest md5_{};
+  char expected_bin_md5_[32];
 };
 
 }  // namespace ota

esphome/components/ota/ota_component.cpp

@@ -8,15 +8,12 @@
 #include "esphome/core/application.h"
 #include "esphome/core/hal.h"
 #include "esphome/core/util.h"
+#include "esphome/components/md5/md5.h"
 #include "esphome/components/network/util.h"
 
 #include <cerrno>
 #include <cstdio>
 
-#ifdef USE_OTA_PASSWORD
-#include <MD5Builder.h>
-#endif
-
 namespace esphome {
 namespace ota {
 
@@ -173,12 +170,12 @@ void OTAComponent::handle_() {
   if (!this->password_.empty()) {
     buf[0] = OTA_RESPONSE_REQUEST_AUTH;
     this->writeall_(buf, 1);
-    MD5Builder md5_builder{};
-    md5_builder.begin();
+    md5::MD5Digest md5{};
+    md5.init();
     sprintf(sbuf, "%08X", random_uint32());
-    md5_builder.add(sbuf);
-    md5_builder.calculate();
-    md5_builder.getChars(sbuf);
+    md5.add(sbuf, 8);
+    md5.calculate();
+    md5.get_hex(sbuf);
     ESP_LOGV(TAG, "Auth: Nonce is %s", sbuf);
 
     // Send nonce, 32 bytes hex MD5
@@ -188,10 +185,10 @@ void OTAComponent::handle_() {
     }
 
     // prepare challenge
-    md5_builder.begin();
-    md5_builder.add(this->password_.c_str());
+    md5.init();
+    md5.add(this->password_.c_str(), this->password_.length());
     // add nonce
-    md5_builder.add(sbuf);
+    md5.add(sbuf, 32);
 
     // Receive cnonce, 32 bytes hex MD5
     if (!this->readall_(buf, 32)) {
@@ -201,11 +198,11 @@ void OTAComponent::handle_() {
     sbuf[32] = '\0';
     ESP_LOGV(TAG, "Auth: CNonce is %s", sbuf);
     // add cnonce
-    md5_builder.add(sbuf);
+    md5.add(sbuf, 32);
 
     // calculate result
-    md5_builder.calculate();
-    md5_builder.getChars(sbuf);
+    md5.calculate();
+    md5.get_hex(sbuf);
     ESP_LOGV(TAG, "Auth: Result is %s", sbuf);
 
     // Receive result, 32 bytes hex MD5