diff --git a/magisk.py b/magisk.py
index f19aac3..9c37b94 100644
--- a/magisk.py
+++ b/magisk.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import argparse
 import hashlib
 import os
 import platform
@@ -150,24 +151,61 @@ class Magisk:
         
         print("==> Magisk installation completed!")
 
-    def run(self):
+    def cleanup(self, keep_downloads=False):
+        """清理临时文件和目录"""
+        print("==> Cleaning up temporary files...")
+        
+        # 清理解压临时目录
+        if os.path.exists(self.extract_to):
+            shutil.rmtree(self.extract_to, ignore_errors=True)
+            print(f"==> Removed temporary extraction directory: {self.extract_to}")
+        
+        # 可选择性清理下载目录
+        if not keep_downloads and os.path.exists(self.download_loc):
+            shutil.rmtree(self.download_loc, ignore_errors=True)
+            print(f"==> Removed download directory: {self.download_loc}")
+        elif keep_downloads:
+            print(f"==> Keeping download directory: {self.download_loc}")
+
+    def run(self, cleanup_after=True, keep_downloads=True):
         """执行完整的构建流程"""
         try:
             self.download()
             self.extract()
             self.copy()
+            success = True
         except Exception as e:
             print(f"Error: {e}")
-            return False
-        return True
+            success = False
+        finally:
+            # 无论成功还是失败都清理临时文件
+            if cleanup_after:
+                self.cleanup(keep_downloads=keep_downloads)
+        
+        return success
 
 
 def main():
     """主函数"""
+    parser = argparse.ArgumentParser(description='Magisk vendor package builder')
+    parser.add_argument('--no-cleanup', action='store_true', 
+                        help='不清理临时文件和下载文件')
+    parser.add_argument('--remove-downloads', action='store_true',
+                        help='同时删除下载的文件')
+    
+    args = parser.parse_args()
+    
+    # 根据参数确定清理行为
+    cleanup_after = not args.no_cleanup
+    keep_downloads = not args.remove_downloads
+    
     magisk = Magisk()
-    success = magisk.run()
+    success = magisk.run(cleanup_after=cleanup_after, keep_downloads=keep_downloads)
+    
     if success:
         print("==> Magisk vendor package created successfully!")
+        if cleanup_after:
+            print("==> Temporary files cleaned up")
     else:
         print("==> Failed to create Magisk vendor package!")
         exit(1)