diff --git a/mocks/generated/trezorconfig.py b/mocks/generated/trezorconfig.py
index a3eb8f7bb..7df40decf 100644
--- a/mocks/generated/trezorconfig.py
+++ b/mocks/generated/trezorconfig.py
@@ -1,3 +1,4 @@
+from typing import *
 
 # extmod/modtrezorconfig/modtrezorconfig.c
 class Config:
diff --git a/mocks/generated/trezorcrypto.py b/mocks/generated/trezorcrypto.py
index c41bfdcb9..1102ab7c1 100644
--- a/mocks/generated/trezorcrypto.py
+++ b/mocks/generated/trezorcrypto.py
@@ -1,3 +1,4 @@
+from typing import *
 
 # extmod/modtrezorcrypto/modtrezorcrypto-aes.h
 class AES:
@@ -5,7 +6,7 @@ class AES:
     AES context.
     '''
 
-    def __init__(self, mode: int, key: bytes, iv: bytes = ...) -> None:
+    def __init__(self, mode: int, key: bytes, iv: bytes = None) -> None:
         '''
         Initialize AES context.
         '''
diff --git a/mocks/generated/trezorio.py b/mocks/generated/trezorio.py
index ed4ff815f..a79f2d101 100644
--- a/mocks/generated/trezorio.py
+++ b/mocks/generated/trezorio.py
@@ -1,3 +1,4 @@
+from typing import *
 
 # extmod/modtrezorio/modtrezorio-sdcard.h
 class SDCard:
diff --git a/mocks/generated/trezormsg.py b/mocks/generated/trezormsg.py
index 4a779058d..ff54ab9b2 100644
--- a/mocks/generated/trezormsg.py
+++ b/mocks/generated/trezormsg.py
@@ -1,3 +1,4 @@
+from typing import *
 
 # extmod/modtrezormsg/modtrezormsg.c
 class HID:
diff --git a/mocks/generated/trezorui.py b/mocks/generated/trezorui.py
index 2e63d82be..e014cf7c6 100644
--- a/mocks/generated/trezorui.py
+++ b/mocks/generated/trezorui.py
@@ -1,3 +1,4 @@
+from typing import *
 
 # extmod/modtrezorui/modtrezorui-display.h
 class Display:
diff --git a/mocks/generated/trezorutils.py b/mocks/generated/trezorutils.py
index 4f6704441..2e09f5398 100644
--- a/mocks/generated/trezorutils.py
+++ b/mocks/generated/trezorutils.py
@@ -1,3 +1,4 @@
+from typing import *
 
 # extmod/modtrezorutils/modtrezorutils.c
 def memcpy(dst: bytearray, dst_ofs: int,
diff --git a/mocks/sys.pyi b/mocks/sys.pyi
index 05b4673b0..f6b91dfc0 100644
--- a/mocks/sys.pyi
+++ b/mocks/sys.pyi
@@ -1,3 +1,5 @@
+from typing import *
+
 def exit(retval: Any = ...) -> None:
     raise SystemExit()
 
diff --git a/mocks/ustruct.pyi b/mocks/ustruct.pyi
index 8791909fe..f61a28ada 100644
--- a/mocks/ustruct.pyi
+++ b/mocks/ustruct.pyi
@@ -1,3 +1,5 @@
+from typing import *
+
 def calcsize(fmt: str) -> int: ...
 def pack(fmt: str, *args: Any) -> bytes: ...
 def pack_into(fmt: str, buffer: bytearray, offset: int, *args: Any) -> None: ...
diff --git a/tools/build_mocks b/tools/build_mocks
index b22893e99..7a4bdb543 100755
--- a/tools/build_mocks
+++ b/tools/build_mocks
@@ -38,15 +38,21 @@ def split_to_parts(line, mod_desc=None):
 
 def store_to_file(dest, parts):
     for package, line in parts:
-        dir_path = os.path.abspath(dest)
+        dirpath = os.path.abspath(dest)
         filename = package
 
-        if not os.path.exists(dir_path):
-            os.makedirs(dir_path)
-            open(os.path.join(dir_path, '__init__.py'), 'w').close()
-            open(os.path.join(dir_path, '.mock-generated'), 'w').close()
+        if not os.path.exists(dirpath):
+            os.makedirs(dirpath)
+            open(os.path.join(dirpath, '__init__.py'), 'w').close()
+            open(os.path.join(dirpath, '.mock-generated'), 'w').close()
 
-        with open(os.path.join(dir_path, filename + '.py'), 'a') as f:
+        filepath = os.path.join(dirpath, filename + '.py')
+
+        if not os.path.exists(filepath):
+            with open(filepath, 'a') as f:
+                f.write('from typing import *\n')
+
+        with open(filepath, 'a') as f:
             f.write(line)