From 62626ac4b98a480bbfd32031fdd9da323f8c7639 Mon Sep 17 00:00:00 2001
From: Piyush Shah <piyush@espressif.com>
Date: Sun, 5 Jan 2020 23:36:48 +0530
Subject: [PATCH] esp_prov: Add support for sending some abstract custom data

To be used with the provisioning/wifi_prov_mgr example

Signed-off-by: Piyush Shah <piyush@espressif.com>
---
 tools/esp_prov/esp_prov.py         | 22 ++++++++++++++++++++++
 tools/esp_prov/prov/custom_prov.py | 14 ++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/tools/esp_prov/esp_prov.py b/tools/esp_prov/esp_prov.py
index 52fb394f2..8ecae3b0d 100644
--- a/tools/esp_prov/esp_prov.py
+++ b/tools/esp_prov/esp_prov.py
@@ -186,6 +186,16 @@ def custom_config(tp, sec, custom_info, custom_ver):
         return None
 
 
+def custom_data(tp, sec, custom_data):
+    try:
+        message = prov.custom_data_request(sec, custom_data)
+        response = tp.send_data('custom-data', message)
+        return (prov.custom_data_response(sec, response) == 0)
+    except RuntimeError as e:
+        on_except(e)
+        return None
+
+
 def scan_wifi_APs(sel_transport, tp, sec):
     APs = []
     group_channels = 0
@@ -328,6 +338,11 @@ if __name__ == '__main__':
                             'If Wi-Fi scanning is supported by the provisioning service, this need not '
                             'be specified'))
 
+    parser.add_argument("--custom_data", dest='custom_data', type=str, default='',
+                        help=desc_format(
+                            'This is an optional parameter, only intended for use with '
+                            '"examples/provisioning/wifi_prov_mgr_custom_data"'))
+
     parser.add_argument("--custom_config", action="store_true",
                         help=desc_format(
                             'This is an optional parameter, only intended for use with '
@@ -394,6 +409,13 @@ if __name__ == '__main__':
             exit(5)
         print("==== Custom config sent successfully ====")
 
+    if args.custom_data != '':
+        print("\n==== Sending Custom data to esp32 ====")
+        if not custom_data(obj_transport, obj_security, args.custom_data):
+            print("---- Error in custom data ----")
+            exit(5)
+        print("==== Custom data sent successfully ====")
+
     if args.ssid == '':
         if not has_capability(obj_transport, 'wifi_scan'):
             print("---- Wi-Fi Scan List is not supported by provisioning service ----")
diff --git a/tools/esp_prov/prov/custom_prov.py b/tools/esp_prov/prov/custom_prov.py
index fe611dbe9..06ec1ca71 100644
--- a/tools/esp_prov/prov/custom_prov.py
+++ b/tools/esp_prov/prov/custom_prov.py
@@ -44,3 +44,17 @@ def custom_config_response(security_ctx, response_data):
     cmd_resp.ParseFromString(decrypt)
     print_verbose(security_ctx, "CustomConfig status " + str(cmd_resp.status))
     return cmd_resp.status
+
+
+def custom_data_request(security_ctx, data):
+    # Encrypt the custom data
+    enc_cmd = security_ctx.encrypt_data(data)
+    print_verbose(security_ctx, "Client -> Device (CustomData cmd) " + utils.str_to_hexstr(enc_cmd))
+    return enc_cmd
+
+
+def custom_data_response(security_ctx, response_data):
+    # Decrypt response packet
+    decrypt = security_ctx.decrypt_data(tobytes(response_data))
+    print("CustomData response: " + str(decrypt))
+    return 0