From 3374d4b8a5a58941e309b13524067ffa8d410d45 Mon Sep 17 00:00:00 2001
From: Sergey Matveev <stargrave@stargrave.org>
Date: Wed, 2 Feb 2022 19:44:53 +0300
Subject: [PATCH] Fix wrap_cryptopro's sbox handling

---
 news.texi           |  4 ++++
 pygost/__init__.py  |  2 +-
 pygost/test_wrap.py | 31 +++++++++++++++++--------------
 pygost/wrap.py      |  7 ++++++-
 4 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/news.texi b/news.texi
index 6ad35a0..96db53f 100644
--- a/news.texi
+++ b/news.texi
@@ -3,6 +3,10 @@
 
 @table @strong
 
+@anchor{Release 5.9}
+@item 5.9
+Fixed @code{wrap.wrap_cryptopro}, that ignored Sbox for key diversification.
+
 @anchor{Release 5.8}
 @item 5.8
 Added human-readable name of the curve in @code{GOST3410Curve.name}.
diff --git a/pygost/__init__.py b/pygost/__init__.py
index 497345c..7800ed4 100644
--- a/pygost/__init__.py
+++ b/pygost/__init__.py
@@ -3,4 +3,4 @@
 PyGOST is free software: see the file COPYING for copying conditions.
 """
 
-__version__ = "5.8"
+__version__ = "5.9"
diff --git a/pygost/test_wrap.py b/pygost/test_wrap.py
index b9d376d..3a0d975 100644
--- a/pygost/test_wrap.py
+++ b/pygost/test_wrap.py
@@ -17,6 +17,7 @@
 from os import urandom
 from unittest import TestCase
 
+from pygost.gost28147 import DEFAULT_SBOX
 from pygost.gost3412 import GOST3412Kuznechik
 from pygost.gost3412 import GOST3412Magma
 from pygost.utils import hexdec
@@ -30,13 +31,14 @@ from pygost.wrap import wrap_gost
 
 class WrapGostTest(TestCase):
     def test_symmetric(self):
-        for _ in range(1 << 8):
-            kek = urandom(32)
-            cek = urandom(32)
-            ukm = urandom(8)
-            wrapped = wrap_gost(ukm, kek, cek)
-            unwrapped = unwrap_gost(kek, wrapped)
-            self.assertSequenceEqual(unwrapped, cek)
+        for sbox in (DEFAULT_SBOX, "id-tc26-gost-28147-param-Z"):
+            for _ in range(1 << 8):
+                kek = urandom(32)
+                cek = urandom(32)
+                ukm = urandom(8)
+                wrapped = wrap_gost(ukm, kek, cek, sbox=sbox)
+                unwrapped = unwrap_gost(kek, wrapped, sbox=sbox)
+                self.assertSequenceEqual(unwrapped, cek)
 
     def test_invalid_length(self):
         with self.assertRaises(ValueError):
@@ -47,13 +49,14 @@ class WrapGostTest(TestCase):
 
 class WrapCryptoproTest(TestCase):
     def test_symmetric(self):
-        for _ in range(1 << 8):
-            kek = urandom(32)
-            cek = urandom(32)
-            ukm = urandom(8)
-            wrapped = wrap_cryptopro(ukm, kek, cek)
-            unwrapped = unwrap_cryptopro(kek, wrapped)
-            self.assertSequenceEqual(unwrapped, cek)
+        for sbox in (DEFAULT_SBOX, "id-tc26-gost-28147-param-Z"):
+            for _ in range(1 << 8):
+                kek = urandom(32)
+                cek = urandom(32)
+                ukm = urandom(8)
+                wrapped = wrap_cryptopro(ukm, kek, cek, sbox=sbox)
+                unwrapped = unwrap_cryptopro(kek, wrapped, sbox=sbox)
+                self.assertSequenceEqual(unwrapped, cek)
 
 
 class TestVectorKExp15(TestCase):
diff --git a/pygost/wrap.py b/pygost/wrap.py
index 9deeff5..eb3855f 100644
--- a/pygost/wrap.py
+++ b/pygost/wrap.py
@@ -79,7 +79,12 @@ def wrap_cryptopro(ukm, kek, cek, sbox=DEFAULT_SBOX):
     :returns: wrapped key
     :rtype: bytes, 44 bytes
     """
-    return wrap_gost(ukm, diversify(kek, bytearray(ukm)), cek, sbox=sbox)
+    return wrap_gost(
+        ukm,
+        diversify(kek, bytearray(ukm), sbox=sbox),
+        cek,
+        sbox=sbox,
+    )
 
 
 def unwrap_cryptopro(kek, data, sbox=DEFAULT_SBOX):
-- 
2.44.0