Dev fixups

.github/workflows/record.yml

@@ -2,59 +2,64 @@ name: Record
 
 on:
   push:
-    branches: [ dev, main, build_exe_2 ]
+    branches: [ dev, main, dev-fixups ]
   workflow_dispatch:
 jobs:
   make-gif:
-    if: contains(github.event.head_commit.message, '[docs]') # only run if commit message contains [docs]
+    if: ${{ contains(github.event.head_commit.message, '[docs]') || github.event_name == 'workflow_dispatch' }} # only run if commit message contains [docs] or workflow_dispatch
     runs-on: ubuntu-latest
+    env:
+      BRANCH_NAME: ${{ github.head_ref || github.ref_name }}
     steps:
     - uses: actions/checkout@v3
-    - name: Set up Python 3.9
+    - name: Set up Python 3.10
       uses: actions/setup-python@v4
       with:
-        python-version: 3.9
+        python-version: "3.10"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install -r requirements.txt
+        pip install termcolor
     - name: Install asciinema
       run: |
         sudo apt-add-repository ppa:zanchey/asciinema -y
         sudo apt-get update -y
         sudo apt-get install asciinema -y
-    - name: Record 
+    - name: Recording
       run: |
-        chmod +x docs/gif_automate.sh
-        asciinema rec terminal.asc -c resources/docs/gif_automate.sh
+        # go through all json in src, put sh into dist
+        python3 resources/actions/autocon/gif_automate.py asc
     - name: Download agg
       run: wget https://github.com/asciinema/agg/releases/latest/download/agg-x86_64-unknown-linux-gnu -O agg
     - name: Make agg executable
       run: chmod +x agg
     - name: Convert to gif
-      run: ./agg terminal.asc terminal.gif
+      run: python3 resources/actions/autocon/gif_automate.py gif
     - name: Upload gif
       uses: actions/upload-artifact@v3
       with:
-        name: terminal.gif
-        path: terminal.gif
+      # upload all gifs in resources/actions/autocon/dist
+        name: gifs
+        path: resources/actions/autocon/dist
     - name: Upload asciinema
       uses: actions/upload-artifact@v3
       with:
-        name: terminal.asc
-        path: terminal.asc
+      # upload all ascs in resources/actions/autocon/src
+        name: ascs
+        path: resources/actions/autocon/src/*.asc
     - name: Move gif
       run: |
         mkdir -p docs
-        mv terminal.gif resources/docs/terminal.gif
+        mv resources/actions/autocon/dist/*.gif resources/docs/
     - name: Commit changes
       run: |
         git config --local user.name "GitHub Actions"
-        git add resources/docs/terminal.gif
-        git commit -m 'Update terminal.gif'
+        git add resources/docs/
+        git commit -m 'Update gifs'
     - name: Push changes
       uses: ad-m/github-push-action@master
       with:
         github_token: ${{ secrets.GITHUB_TOKEN }}
-        branch: build_exe_2
+        branch: ${{ env.BRANCH_NAME }}
 

.gitignore

@@ -135,3 +135,5 @@ aimodels-lock.json
 /ldm/
 /.fleet/
 
+# For local aimodels packages
+aimodels-lock.json

aimm.py

@@ -5,7 +5,7 @@
 WEBSITE = "https://aimodels.org"
 API_SERVER = "https://api.aimodels.org"
 GITHUB_REPO = "visioninit/aimm"
-VERSION = "0.3.0"
+VERSION = "alpha-1"
 
 
 def init():

cli/__init__.py

@@ -1,4 +1,4 @@
-# initialize the info command
+# initialize the commands
 from cli import (base_funcs,
                  install,
                  uninstall,
@@ -9,5 +9,6 @@
                  add,
                  remove,
                  security,
-                 credentials
+                 credentials,
+                 pickle
                  )

cli/add.py

@@ -8,7 +8,7 @@
 
 app = aimmApp.app
 @app.command()
-def add(name_version: str):
+def add(name_version: str, mut_path: bool = typer.Option(False, "--allow-mutable-paths")):
     """
     Add a model to local aimodels.json.
     """
@@ -22,9 +22,12 @@ def add(name_version: str):
         sys.exit(1)
 
     if base_funcs.should_install(name, version):
-        install.install(name_version)
+        install.install(name_version,mut_path=mut_path)
     else:
-        typer.echo(f"{name}:{version} already installed.")
+        # add to aimodels-lock.json
+        save_path = os.path.join(aimmApp.main_dir, name, version)
+        base_funcs.update_ai_models_lock(name, version, save_path)
+        typer.echo(f"Is this used: {name}:{version}")
 
     # parse aimodels.json as a json
     try:

cli/aimmApp.py

@@ -56,6 +56,7 @@ def show_help(program_name):
     print("  info <model_name>:[version]                 Get info about a model (name, version, description, etc).")
     print("  install <model_name>:[version]              Install a model to system.")
     print("  uninstall <model_name>:[version]            Uninstall a model.")
+    print("  scan <model_name>:[version]                 Scan model for pickles and potential issues.")
     print()
     print("Search:")
     print("  search <query>                              Search for a model.")
@@ -77,7 +78,7 @@ def check_for_updates(repo, current_version):
     try:
         r = requests.get(url, timeout=10)
         if r.status_code == 200:
-            latest_version = r.json()[0]["tag_name"]
+            latest_version = r.json()["tag_name"]
             latest_version_list = latest_version.split(".")
             current_version_list = current_version.split(".")
             if (latest_version_list[0] > current_version_list[0]) or (latest_version_list[1] > current_version_list[1]) or (latest_version_list[2] > current_version_list[2]):

cli/base_funcs.py

@@ -127,7 +127,8 @@ def update_ai_models_lock(name, version, path):
     aimodels_lock_file = os.path.join(cwd, "aimodels-lock.json")
     # get a list of files in path
     files = os.listdir(path)
-
+    # path should use forward slashes instead of backslashes
+    path = path.replace("\\", "/")
     # check if aimodels-lock.json exists
     if not os.path.exists(aimodels_lock_file):
         # create aimodels-lock.json
@@ -137,8 +138,8 @@ def update_ai_models_lock(name, version, path):
                     "path": path,
                     "files": files
                 },
+            },
             "credentials": {}
-            }
         }
         # write aimodels-lock.json
         with open("aimodels-lock.json", "w") as f:
@@ -183,7 +184,7 @@ def update_ai_models_lock(name, version, path):
     else:
         with open(".gitignore", "r") as f:
             gitignore = f.read()
-        if gitignore_text not in gitignore:
+        if "aimodels-lock.json" not in gitignore:
             # append to end
             with open(".gitignore", "a") as f:
                 f.write(gitignore_text)
@@ -212,4 +213,15 @@ def get_creds(download_url):
         with open("aimodels-lock.json", "r") as f:
             aimodels_lock = json.load(f)
         # return username and password
-        return aimodels_lock["credentials"][domain]["username"], aimodels_lock["credentials"][domain]["password"]
+        return aimodels_lock["credentials"][domain]["username"], aimodels_lock["credentials"][domain]["password"]
+
+def get_model_path(name_version):
+    name, version = extract_name_version(name_version)
+    if version is None:
+        version = get_last_version(name)
+    for package in aimmApp.installed["packages"]:
+        if package["name"].lower() == name.lower() and package["version"] == version:
+            return package["paths"]
+    else:
+        typer.echo(f"Error: {name}:{version} not found")
+        return None

cli/install.py

@@ -42,7 +42,7 @@ def install(name_version: Optional[str] = typer.Argument(None),
         return
 
     if base_funcs.should_install(name, version):
-        typer.echo(f"Installing {name}:{version}...")
+        typer.echo(f"Install Pre-Check: {name}:{version}...")
         # download the model
         # if data doesn't have entries exit
         if base_funcs.models_json(name,version)["data"] == []:
@@ -61,13 +61,23 @@ def install(name_version: Optional[str] = typer.Argument(None),
                 if not file["download_url"]:
                     typer.echo(f"Error: Model {name}:{version} not found")
                     return
-                # check if md5 exists
-                # check if mut_path is passed 
+                # check if mut_path is passed
                 if not mut_path:
-                    if not file["md5"]:
-                        typer.echo(f"Error: {name}:{version}")
-                        typer.echo("  Remote URL file could be changed and no checksum provided to validate file.")
-                        sys.exit(1)
+                    typer.echo(f"Error: Unverifiable download path in request")
+                    typer.echo("  * Remote file is mutable\n"
+                               "    Immutable download paths always point to the same file. In cases\n"
+                               "    where the file is mutable, the file could be changed.\n\n"
+                               "    Example of immutable paths:\n"
+                               "    - github.com/{user}}/{repo_name}/tree/{commit_hash}\n" 
+                               "    - huggingface.co/spaces/{space_name}/{repo_name}/tree/{commit_hash}\n\n" 
+                               "    Mutable paths in your request:\n"
+                               "    - somemodel:version\n"
+                               "      "+file['download_url']+"\n"
+                               "    -  othermodel:version\n"
+                               "      "+file['download_url']+"\n")
+                    typer.echo("   To allow mutable files, run command again with argument:\n"
+                               "\t --allow-mutable-paths")
+                    sys.exit(1)
 
                 save_path = os.path.join(aimmApp.main_dir, name, version)
                 if not os.path.exists(save_path):
@@ -88,8 +98,8 @@ def install(name_version: Optional[str] = typer.Argument(None),
                 else:
                     base_funcs.download_file(file["download_url"], save_path)
 
-                # add to aimodels-lock.json
-                base_funcs.update_ai_models_lock(name, version, save_path)
+            # add to aimodels-lock.json
+            base_funcs.update_ai_models_lock(name, version, save_path)
             # make a list of links
             links = []
             # only execute if there are links
@@ -109,4 +119,6 @@ def install(name_version: Optional[str] = typer.Argument(None),
                 json.dump(aimmApp.installed, file, indent=4)
             typer.echo(f"Installed {name}:{version}!")
     else: 
-        typer.echo(f"{name}:{version} already installed.")
+        save_path = os.path.join(aimmApp.main_dir, name, version)
+        base_funcs.update_ai_models_lock(name, version, save_path)
+        typer.echo(f"Found Local: {name}:{version}")

cli/pickle.py

@@ -0,0 +1,68 @@
+import os
+import hashlib
+import json
+import typer
+from picklescan import scanner
+
+from cli import aimmApp, base_funcs
+
+app = aimmApp.app
+
+def hash_file(filename):
+    hash_obj = hashlib.sha256()
+    with open(filename, 'rb') as file:
+        chunk = 0
+        while chunk != b'':
+            chunk = file.read(1024)
+            hash_obj.update(chunk)
+
+    return hash_obj.hexdigest()
+
+
+def get_maximum_danger(result_globals) -> str:
+    safety = None
+    for module in result_globals:
+        if safety is None:
+            safety = module.safety
+        elif safety.value < module.safety.value:
+            safety = module.safety
+    match safety:
+        case scanner.SafetyLevel.Innocuous:
+            safety = "innocuous"
+        case scanner.SafetyLevel.Suspicious:
+            safety = "suspicious"
+        case scanner.SafetyLevel.Dangerous:
+            safety = "dangerous"
+    return safety
+
+@app.command()
+def scan(name_version: str, raw: bool = typer.Option(False, "--raw", "-r")):
+    """
+    Scan a file for malicious code.
+    """
+    model_path = base_funcs.get_model_path(name_version)
+    if model_path:
+        for file in os.listdir(model_path):
+            path = os.path.join(model_path, file)
+            result = scanner.scan_file_path(path)
+            if not raw:
+                print(f"Scanning {path}...")
+                print(f"scanned files: {result.scanned_files}")
+                print(f"issues count: {result.issues_count}")
+                print(f"infected files: {result.infected_files}")
+                print("safety level of modules:")
+                for module in result.globals:
+                    print(
+                        f"  * {module.module}.{module.name} - {(module.safety.value).title()}")
+                print(f"Hash: {hash_file(path)}")
+            else:
+                safety = get_maximum_danger(result.globals)
+                print(json.dumps({
+                    "hash": hash_file(path),
+                    "picklescan": {
+                        "scanned_files": result.scanned_files,
+                        "issues_count": result.issues_count,
+                        "infected_files": result.infected_files,
+                        "safety_level": safety
+                    }
+                },indent=4))

cli/remove.py

@@ -36,7 +36,7 @@ def remove(name_version: str):
         if package_name.lower() == name.lower() and package_version.lower() == version.lower():
             aimodels.pop(package_name)
             typer.echo(f"Removed {name_version} from aimodels.json")
-            typer.echo(f"The package is still available system-wide, to uninstall it run aimm uninstall {name}:{version}")
+            typer.echo(f"The package is still available system-wide, to uninstall (delete files):\n\t aimm uninstall {name}:{version}")
             break
     else:
         typer.echo(f"{name_version} not found in aimodels.json")

requirements.txt

@@ -11,4 +11,5 @@ tqdm==4.57.0
 typer==0.7.0
 urllib3==1.26.12
 wcwidth==0.2.5
-pip-licenses==4.0.1
+pip-licenses==4.0.1
+picklescan==0.0.5

resources/actions/autocon/gif_automate.py

@@ -0,0 +1,25 @@
+import os
+import sys
+
+DEFAULT_COMMAND = "asciinema rec"
+DEFAULT_DIR = "resources/actions/autocon"
+
+if __name__ == '__main__':
+    if len(sys.argv) >= 2:
+        # change working directory to the main directory
+        os.chdir(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
+        # for every json file run json_typing.py
+        if sys.argv[1] == "asc":
+            for file in os.listdir(f"{DEFAULT_DIR}/src"):
+                if file.endswith('.json'):
+                    name = file.split('.json')[0]
+                    os.system(f'{DEFAULT_COMMAND} {DEFAULT_DIR}/src/{name}.asc -c "python3 {DEFAULT_DIR}/json_typing.py {DEFAULT_DIR}/src/{name}.json"')
+        elif sys.argv[1] == "gif":
+            # make dist folder
+            if not os.path.exists(f"{DEFAULT_DIR}/dist"):
+                os.mkdir(f"{DEFAULT_DIR}/dist")
+            for file in os.listdir(f"{DEFAULT_DIR}/src"):
+                if file.endswith('.asc'):
+                    name = file.split('.asc')[0]
+                    os.system(f'./agg {DEFAULT_DIR}/src/{name}.asc {DEFAULT_DIR}/dist/{name}.gif')
+