Split install tools and packages

Signed-off-by: Huaqi Fang <578567190@qq.com>
Nuclei-Software · Dec 29, 2023 · 7889b30 · 7889b30
1 parent b29155a
commit 7889b30
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 12 deletions.
diff --git a/.github/prepare_tools.py b/.github/prepare_tools.py
@@ -5,14 +5,29 @@
 import shutil
 import json
 import requests
-import wget  # Import the wget library
 from urllib.parse import urlparse
 import argparse
+try:
+    import wget  # Import the wget library
+    NOWGET = False
+except:
+    NOWGET = True
+    pass
 
 PIOJSONLOC = ""
 PREBLT_CACHE = "prebuilt_dlcache"
 PREBLT_TOOLS = "prebuilt_tools"
 
+def download_file(url, file_name):
+    if NOWGET:
+        # Download the file using stream to avoid loading the entire file into memory
+        with requests.get(url, stream=True) as response:
+            with open(file_name, 'wb') as file:
+                shutil.copyfileobj(response.raw, file)
+    else:
+        wget.download(url, file_name)
+    pass
+
 def download_with_progress(url, destination_folder, reuse=False):
     file_name = os.path.join(destination_folder, os.path.basename(urlparse(url).path))
     if os.path.isdir(destination_folder) == False:
@@ -22,16 +37,11 @@ def download_with_progress(url, destination_folder, reuse=False):
         os.remove(file_name)
     # Download the file with progress bar
     if os.path.isfile(file_name) == False:
-        wget.download(url, file_name)
+        download_file(url, file_name)
         print("%s is downloaded!" % (file_name))
     else:
         print("%s already downloaded!" % (file_name))
 
-    # Download the file using stream to avoid loading the entire file into memory
-    #with requests.get(url, stream=True) as response:
-    #    with open(file_name, 'wb') as file:
-    #        shutil.copyfileobj(response.raw, file)
-
     return file_name
 
 def download_and_extract(url, extract_folder, reuse=False):
@@ -127,7 +137,6 @@ def prepare_tools():
     pass
 
 def install_pio_packages(nsdk_url="https://github.com/Nuclei-Software/nuclei-sdk#feature/gd32vw55x"):
-    print("Install pio required packages")
     if os.path.isfile("platform.py") == False:
         print("Not in platform nuclei folder, exit")
         return
@@ -144,14 +153,18 @@ def install_pio_packages(nsdk_url="https://github.com/Nuclei-Software/nuclei-sdk
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Prepare Nuclei Tools and Install PIO Packages.')
     parser.add_argument('--sdk', "-s", default="https://github.com/Nuclei-Software/nuclei-sdk#feature/gd32vw55x", help='URL or PATH of Nuclei SDK')
+    parser.add_argument('--pio', action='store_true', help="Setup PIO Package")
 
     args = parser.parse_args()
 
     if os.path.isdir(PREBLT_TOOLS):
         print("%s existed, maybe tools already installed!" % (PREBLT_TOOLS))
     else:
         prepare_tools()
-
-    install_pio_packages(args.sdk)
+
+    if args.pio:
+        print("Install pio required packages")
+        install_pio_packages(args.sdk)
 
     print("Setup completed successfully!")
+    sys.exit(0)
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -29,9 +29,9 @@ jobs:
           path: prebuilt_dlcache
           key: build
 
-      - name: Install Nuclei Tools and PIO Packages
+      - name: Install and Setup Nuclei Tools
         run: |
-          python3 .github/prepare_tools.py
+          python3 -u .github/prepare_tools.py
 
       - name: Upload cached packages
         uses: actions/upload-artifact@v3
@@ -40,6 +40,10 @@ jobs:
           path: |
             prebuilt_dlcache
 
+      - name: Setup PIO Packages
+        run: |
+          python3 -u .github/prepare_tools.py --pio
+
       - name: Build examples
         run: |
           python3 .github/build_examples.py