From afd54d11d3a663300e9e895aacdbc16ecf914a9a Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 29 Jun 2023 19:43:01 -0500
Subject: [PATCH] Add optional kasa-crypt dependency for speedups (#464)

If installed, use the optimized protocol encryption procedures implemented as a C extension by kasa-crypt (https://pypi.org/project/kasa-crypt/
---
 .github/workflows/ci.yml | 28 ++++++++++++++++++++++++++--
 kasa/protocol.py         | 10 ++++++++++
 poetry.lock              | 15 ++++++++++++++-
 pyproject.toml           |  4 +++-
 4 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index b5a3b310..d0dce418 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -58,7 +58,7 @@ jobs:
           poetry run pre-commit run check-ast --all-files
 
   tests:
-    name: "Python ${{ matrix.python-version}} on ${{ matrix.os }}"
+    name: Python ${{ matrix.python-version}} on ${{ matrix.os }}${{ fromJSON('[" (extras)", ""]')[matrix.extras == ''] }}
     needs: linting
     runs-on: ${{ matrix.os }}
 
@@ -66,6 +66,24 @@ jobs:
       matrix:
         python-version: ["3.8", "3.9", "3.10", "3.11", "pypy-3.8"]
         os: [ubuntu-latest, macos-latest, windows-latest]
+        extras: [false, true]
+        exclude:
+          - os: macos-latest
+            extras: true
+          - os: windows-latest
+            extras: true
+          - os: ubuntu-latest
+            python-version: "pypy-3.8"
+            extras: true
+          - os: ubuntu-latest
+            python-version: "3.8"
+            extras: true
+          - os: ubuntu-latest
+            python-version: "3.9"
+            extras: true
+          - os: ubuntu-latest
+            python-version: "3.10"
+            extras: true
         # exclude pypy on windows, as the poetry install seems to be very flaky:
         # PermissionError(13, 'The process cannot access the file because it is being used by another process'))
         # at C:\hostedtoolcache\windows\PyPy\3.7.10\x86\site-packages\requests\models.py:761 in generate
@@ -78,10 +96,16 @@ jobs:
       - uses: "actions/setup-python@v2"
         with:
           python-version: "${{ matrix.python-version }}"
-      - name: "Install dependencies"
+      - name: "Install dependencies (no speedups)"
+        if: matrix.extras == false
         run: |
           python -m pip install --upgrade pip poetry
           poetry install
+      - name: "Install dependencies (with speedups)"
+        if: matrix.extras == true
+        run: |
+          python -m pip install --upgrade pip poetry
+          poetry install --extras speedups
       - name: "Run tests"
         run: |
           poetry run pytest --cov kasa --cov-report xml
diff --git a/kasa/protocol.py b/kasa/protocol.py
index bcb8b7f5..e21866e3 100755
--- a/kasa/protocol.py
+++ b/kasa/protocol.py
@@ -219,3 +219,13 @@ class TPLinkSmartHomeProtocol:
         return bytes(
             TPLinkSmartHomeProtocol._xor_encrypted_payload(ciphertext)
         ).decode()
+
+
+# Try to load the kasa_crypt module and if it is available
+try:
+    from kasa_crypt import decrypt, encrypt
+
+    TPLinkSmartHomeProtocol.decrypt = decrypt  # type: ignore[method-assign]
+    TPLinkSmartHomeProtocol.encrypt = encrypt  # type: ignore[method-assign]
+except ImportError:
+    pass
diff --git a/poetry.lock b/poetry.lock
index 065c6449..57149e11 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -445,6 +445,18 @@ MarkupSafe = ">=2.0"
 [package.extras]
 i18n = ["Babel (>=2.7)"]
 
+[[package]]
+name = "kasa-crypt"
+version = "0.2.0"
+description = "Fast kasa crypt"
+category = "main"
+optional = true
+python-versions = ">=3.7,<4.0"
+files = [
+    {file = "kasa_crypt-0.2.0-cp310-cp310-manylinux_2_31_x86_64.whl", hash = "sha256:9676ac702aa62252fb4de3d5c9ee4895dd93610b9c37e732213b0914fbc0e255"},
+    {file = "kasa_crypt-0.2.0.tar.gz", hash = "sha256:ad2e73276f09ed035d53006985b08eb78869f73e60ac5d66547d9ddc35cb8cc4"},
+]
+
 [[package]]
 name = "markdown-it-py"
 version = "2.2.0"
@@ -1387,8 +1399,9 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more
 
 [extras]
 docs = ["sphinx", "sphinx_rtd_theme", "sphinxcontrib-programoutput", "myst-parser", "docutils"]
+speedups = ["orjson", "kasa-crypt"]
 
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8"
-content-hash = "ea746265dddc60d13ef0edb476d02181631519af39fa90574e60f8eabe35bbe0"
+content-hash = "b85f55f0ca928b1f3510da37196c21f40eb07cd4d07b3a9c3dd29215ba9777fe"
diff --git a/pyproject.toml b/pyproject.toml
index 62b8f4b0..57c1f21c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,8 @@ asyncclick = ">=8"
 pydantic = "^1"
 
 # speed ups
-orjson = { "version" = ">=3.9.1", optional = true, extras = ["speedups"] }
+orjson = { "version" = ">=3.9.1", optional = true }
+kasa-crypt = { "version" = ">=0.2.0", optional = true }
 
 # required only for docs
 sphinx = { version = "^4", optional = true }
@@ -50,6 +51,7 @@ coverage = {version = "*", extras = ["toml"]}
 
 [tool.poetry.extras]
 docs = ["sphinx", "sphinx_rtd_theme", "sphinxcontrib-programoutput", "myst-parser", "docutils"]
+speedups = ["orjson", "kasa-crypt"]
 
 
 [tool.isort]