python: use stdlib blake2s on python 3.6+

python/requirements.txt

@@ -3,7 +3,7 @@ mnemonic>=0.17
 shamir-mnemonic>=0.1.0
 requests>=2.4.0
 click>=7,<8
-pyblake2>=0.9.3
 libusb1>=1.6.4
 construct>=2.9
 typing_extensions>=3.6
+pyblake2>=0.9.3 ; python_version<'3.6'

python/setup.py

@@ -12,10 +12,10 @@ install_requires = [
     "shamir-mnemonic>=0.1.0",
     "requests>=2.4.0",
     "click>=7,<8",
-    "pyblake2>=0.9.3",
     "libusb1>=1.6.4",
     "construct>=2.9",
     "typing_extensions>=3.6",
+    "pyblake2>=0.9.3 ; python_version<'3.6'",
 ]
 
 CWD = os.path.dirname(os.path.realpath(__file__))

python/src/trezorlib/firmware.py

@@ -20,10 +20,15 @@ from typing import Callable, List, NewType, Tuple
 
 import construct as c
 import ecdsa
-import pyblake2
 
 from . import cosi, messages, tools
 
+try:
+    from hashlib import blake2s
+except ImportError:
+    from pyblake2 import blake2s
+
+
 V1_SIGNATURE_SLOTS = 3
 V1_BOOTLOADER_KEYS = {
     1: "04d571b7f148c5e4232c3814f777d8faeaf1a84216c78d569b71041ffc768a5b2d810fc3bb134dd026b57e65005275aedef43e155f48fc11a32ec790a93312bd58",
@@ -263,9 +268,7 @@ def check_sig_v1(
 
 
 def _header_digest(
-    header: c.Container,
+    header: c.Container, header_type: c.Construct, hash_function: Callable = blake2s
-    header_type: c.Construct,
-    hash_function: Callable = pyblake2.blake2s,
 ) -> bytes:
     stripped_header = header.copy()
     stripped_header.sigmask = 0
@@ -277,7 +280,7 @@ def _header_digest(
 
 
 def digest_v2(fw: FirmwareType) -> bytes:
-    return _header_digest(fw.firmware_header, FirmwareHeader, pyblake2.blake2s)
+    return _header_digest(fw.firmware_header, FirmwareHeader, blake2s)
 
 
 def digest_onev2(fw: FirmwareType) -> bytes:
@@ -286,7 +289,7 @@ def digest_onev2(fw: FirmwareType) -> bytes:
 
 def validate_code_hashes(
     fw: FirmwareType,
-    hash_function: Callable = pyblake2.blake2s,
+    hash_function: Callable = blake2s,
     chunk_size: int = V2_CHUNK_SIZE,
     padding_byte: bytes = None,
 ) -> None:
@@ -426,7 +429,7 @@ def update(client, data):
     # TREZORv2 method
     while isinstance(resp, messages.FirmwareRequest):
         payload = data[resp.offset : resp.offset + resp.length]
-        digest = pyblake2.blake2s(payload).digest()
+        digest = blake2s(payload).digest()
         resp = client.call(messages.FirmwareUpload(payload=payload, hash=digest))
 
     if isinstance(resp, messages.Success):