diff --git a/tools/build_mocks b/tools/build_mocks
index 4c064a2e8..b22893e99 100755
--- a/tools/build_mocks
+++ b/tools/build_mocks
@@ -3,53 +3,71 @@ import os
 
 COMMENT_PREFIX = '/// '
 
+current_indent = 0
+current_class = None
 current_method = None
 current_package = None
 
 
 def split_to_parts(line, mod_desc=None):
+    global current_indent
+    global current_class
     global current_method
-    global current_package
-    # Line is beginning of method
-    if line.startswith('def '):
-        # Parse name of method from line like 'def trezor.config.get():'
-        current_method = line[4:].split('(')[0]
-        #print("Current method", current_method)
 
-        *current_package, method_name = current_method.split('.')
+    if line.startswith('class '):
+        current_class = line[6:].split('(')[0].strip(':')
+        current_indent = 0
 
         yield (current_package, "\n")
         yield (current_package, '# ' + mod_desc + "\n")
-        line = line.replace(current_method, method_name)
+
+    elif line.startswith('def '):
+        current_method = line[4:].split('(')[0]
+
+        yield (current_package, "\n")
+
+        if current_class is None:
+            yield (current_package, '# ' + mod_desc + "\n")
+        else:
+            current_indent = 4
+
+    line = current_indent * ' ' + line
 
     yield (current_package, line)
 
 
 def store_to_file(dest, parts):
-    while True:
-        try:
-            (package, line) = parts.__next__()
-        except StopIteration:
-            return
-
-        dir_path = os.path.abspath(os.path.join(dest, *package[:-1]))
-        filename = package[-1]
+    for package, line in parts:
+        dir_path = 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()
 
-        f = open(os.path.join(dir_path, filename + '.py'), 'a')
-        f.write(line)
-        f.close()
+        with open(os.path.join(dir_path, filename + '.py'), 'a') as f:
+            f.write(line)
 
 
 def build_module(mod_file, dest):
+    global current_indent
+    global current_class
+    global current_package
+
     if not (mod_file.endswith('.h') or mod_file.endswith('.c')):
         return
+    if not os.path.basename(mod_file).startswith('mod'):
+        return
 
+    current_indent = 0
+    current_class = None
+    current_package = os.path.basename(mod_file) \
+        .split('.')[0] \
+        .split('-')[0] \
+        .replace('mod', '')
     mod_desc = mod_file.replace('../micropython/extmod', 'extmod')
+
     for l in open(mod_file):
         if not l.startswith(COMMENT_PREFIX):
             continue
@@ -69,13 +87,13 @@ def clear_directory(top_dir):
     print("Clearing up directory", top_dir)
     for root, dirs, files in os.walk(top_dir, topdown=False):
         if '.mock-generated' not in files:
-            #print("Not a mock directory", root)
+            # print("Not a mock directory", root)
             continue
         for name in files:
-            #print('Deleting file', os.path.join(root, name))
+            # print('Deleting file', os.path.join(root, name))
             os.remove(os.path.join(root, name))
         for name in dirs:
-            #print('Deleting directory', os.path.join(root, name))
+            # print('Deleting directory', os.path.join(root, name))
             try:
                 os.rmdir(os.path.join(root, name))
             except FileNotFoundError:
@@ -84,5 +102,5 @@ def clear_directory(top_dir):
         os.rmdir(root)
 
 if __name__ == '__main__':
-    clear_directory('../mocks')
-    build_directory('../micropython/extmod', '../mocks')
+    clear_directory('../mocks/generated')
+    build_directory('../micropython/extmod', '../mocks/generated')