From f373c7b93f58285ed063f8af4e65d44b516f6bff Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Tue, 7 Dec 2021 17:11:59 +0200
Subject: [PATCH] tests: DPP PKEXv2

Signed-off-by: Jouni Malinen <quic_jouni@quicinc.com>
---
 tests/hwsim/hostapd.py       |  7 +++++--
 tests/hwsim/test_dpp.py      | 37 ++++++++++++++++++++++++++++++++++--
 tests/hwsim/wpasupplicant.py |  7 +++++--
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/tests/hwsim/hostapd.py b/tests/hwsim/hostapd.py
index 5ea68444c..f9becfcef 100644
--- a/tests/hwsim/hostapd.py
+++ b/tests/hwsim/hostapd.py
@@ -508,7 +508,7 @@ class Hostapd:
             raise Exception("Failed to initiate DPP Authentication")
 
     def dpp_pkex_init(self, identifier, code, role=None, key=None, curve=None,
-                      extra=None, use_id=None):
+                      extra=None, use_id=None, v2=False):
         if use_id is None:
             id1 = self.dpp_bootstrap_gen(type="pkex", key=key, curve=curve)
         else:
@@ -516,7 +516,10 @@ class Hostapd:
         cmd = "own=%d " % id1
         if identifier:
             cmd += "identifier=%s " % identifier
-        cmd += "init=1 "
+        if v2:
+            cmd += "init=2 "
+        else:
+            cmd += "init=1 "
         if role:
             cmd += "role=%s " % role
         if extra:
diff --git a/tests/hwsim/test_dpp.py b/tests/hwsim/test_dpp.py
index ca6092a69..3d9965726 100644
--- a/tests/hwsim/test_dpp.py
+++ b/tests/hwsim/test_dpp.py
@@ -2268,6 +2268,10 @@ def test_dpp_pkex(dev, apdev):
     """DPP and PKEX"""
     run_dpp_pkex(dev, apdev)
 
+def test_dpp_pkex_v2(dev, apdev):
+    """DPP and PKEXv2"""
+    run_dpp_pkex(dev, apdev, v2=True)
+
 def test_dpp_pkex_p256(dev, apdev):
     """DPP and PKEX (P-256)"""
     run_dpp_pkex(dev, apdev, "P-256")
@@ -2321,13 +2325,13 @@ def test_dpp_pkex_identifier_mismatch3(dev, apdev):
 
 def run_dpp_pkex(dev, apdev, curve=None, init_extra=None, check_config=False,
                  identifier_i="test", identifier_r="test",
-                 expect_no_resp=False):
+                 expect_no_resp=False, v2=False):
     check_dpp_capab(dev[0], curve and "brainpool" in curve)
     check_dpp_capab(dev[1], curve and "brainpool" in curve)
     dev[0].dpp_pkex_resp(2437, identifier=identifier_r, code="secret",
                          curve=curve)
     dev[1].dpp_pkex_init(identifier=identifier_i, code="secret", curve=curve,
-                         extra=init_extra)
+                         extra=init_extra, v2=v2)
 
     if expect_no_resp:
         ev = dev[0].wait_event(["DPP-RX"], timeout=10)
@@ -2551,6 +2555,19 @@ def test_dpp_pkex_hostapd_responder(dev, apdev):
     wait_auth_success(hapd, dev[0], configurator=dev[0], enrollee=hapd,
                       stop_initiator=True)
 
+def test_dpp_pkex_v2_hostapd_responder(dev, apdev):
+    """DPP PKEXv2 with hostapd as responder"""
+    check_dpp_capab(dev[0])
+    hapd = hostapd.add_ap(apdev[0], {"ssid": "unconfigured",
+                                     "channel": "6"})
+    check_dpp_capab(hapd)
+    hapd.dpp_pkex_resp(2437, identifier="test", code="secret")
+    conf_id = dev[0].dpp_configurator_add()
+    dev[0].dpp_pkex_init(identifier="test", code="secret",
+                         extra="conf=ap-dpp configurator=%d" % conf_id, v2=True)
+    wait_auth_success(hapd, dev[0], configurator=dev[0], enrollee=hapd,
+                      stop_initiator=True)
+
 def test_dpp_pkex_hostapd_initiator(dev, apdev):
     """DPP PKEX with hostapd as initiator"""
     check_dpp_capab(dev[0])
@@ -2566,6 +2583,22 @@ def test_dpp_pkex_hostapd_initiator(dev, apdev):
     wait_auth_success(hapd, dev[0], configurator=dev[0], enrollee=hapd,
                       stop_initiator=True)
 
+def test_dpp_pkex_v2_hostapd_initiator(dev, apdev):
+    """DPP PKEXv2 with hostapd as initiator"""
+    check_dpp_capab(dev[0])
+    hapd = hostapd.add_ap(apdev[0], {"ssid": "unconfigured",
+                                     "channel": "6"})
+    check_dpp_capab(hapd)
+    conf_id = dev[0].dpp_configurator_add()
+    dev[0].set("dpp_configurator_params",
+               " conf=ap-dpp configurator=%d" % conf_id)
+    dev[0].dpp_pkex_resp(2437, identifier="test", code="secret",
+                         listen_role="configurator")
+    hapd.dpp_pkex_init(identifier="test", code="secret", role="enrollee",
+                       v2=True)
+    wait_auth_success(hapd, dev[0], configurator=dev[0], enrollee=hapd,
+                      stop_initiator=True)
+
 def test_dpp_pkex_hostapd_errors(dev, apdev):
     """DPP PKEX errors with hostapd"""
     hapd = hostapd.add_ap(apdev[0], {"ssid": "unconfigured",
diff --git a/tests/hwsim/wpasupplicant.py b/tests/hwsim/wpasupplicant.py
index fdece92f6..160aa3e2d 100644
--- a/tests/hwsim/wpasupplicant.py
+++ b/tests/hwsim/wpasupplicant.py
@@ -1577,7 +1577,7 @@ class WpaSupplicant:
         return int(peer)
 
     def dpp_pkex_init(self, identifier, code, role=None, key=None, curve=None,
-                      extra=None, use_id=None, allow_fail=False):
+                      extra=None, use_id=None, allow_fail=False, v2=False):
         if use_id is None:
             id1 = self.dpp_bootstrap_gen(type="pkex", key=key, curve=curve)
         else:
@@ -1585,7 +1585,10 @@ class WpaSupplicant:
         cmd = "own=%d " % id1
         if identifier:
             cmd += "identifier=%s " % identifier
-        cmd += "init=1 "
+        if v2:
+            cmd += "init=2 "
+        else:
+            cmd += "init=1 "
         if role:
             cmd += "role=%s " % role
         if extra: