diff --git a/token_extractor.py b/token_extractor.py
index d9f51ac..3c5f69b 100644
--- a/token_extractor.py
+++ b/token_extractor.py
@@ -265,7 +265,14 @@ def main():
     print()
     if not server == "":
         servers = [server]
+    run(username, password, servers)
+    
+    print()
+    print("Press ENTER to finish")
+    input()
+
 
+def run(username, password, servers):
     connector = XiaomiCloudConnector(username, password)
     print("Logging in...")
     logged = connector.login()
@@ -319,10 +326,6 @@ def main():
     else:
         print("Unable to log in.")
 
-    print()
-    print("Press ENTER to finish")
-    input()
-
 
 if __name__ == "__main__":
     main()