diff --git a/scripts/mbedtls_dev/build_tree.py b/scripts/mbedtls_dev/build_tree.py
index f52b785..8640857 100644
--- a/scripts/mbedtls_dev/build_tree.py
+++ b/scripts/mbedtls_dev/build_tree.py
@@ -19,12 +19,19 @@
 import os
 import inspect
 
+def looks_like_psa_crypto_root(path: str) -> bool:
+    """Whether the given directory looks like the root of the PSA Crypto source tree."""
+    return all(os.path.isdir(os.path.join(path, subdir))
+               for subdir in ['include', 'core', 'tests'])
 
 def looks_like_mbedtls_root(path: str) -> bool:
     """Whether the given directory looks like the root of the Mbed TLS source tree."""
     return all(os.path.isdir(os.path.join(path, subdir))
                for subdir in ['include', 'library', 'programs', 'tests'])
 
+def looks_like_root(path: str) -> bool:
+    return looks_like_psa_crypto_root(path) or looks_like_mbedtls_root(path)
+
 def check_repo_path():
     """
     Check that the current working directory is the project root, and throw
@@ -42,7 +49,7 @@
     for d in [os.path.curdir,
               os.path.pardir,
               os.path.join(os.path.pardir, os.path.pardir)]:
-        if looks_like_mbedtls_root(d):
+        if looks_like_root(d):
             os.chdir(d)
             return
     raise Exception('Mbed TLS source tree not found')
@@ -62,6 +69,6 @@
             if d in dirs:
                 continue
             dirs.add(d)
-            if looks_like_mbedtls_root(d):
+            if looks_like_root(d):
                 return d
     raise Exception('Mbed TLS source tree not found')