init scan

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,14 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="Eslint" enabled="true" level="WARNING" enabled_by_default="true" />
+    <inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="N806" />
+          <option value="N802" />
+        </list>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="$USER_HOME$/miniconda3" project-jdk-type="Python SDK" />
+</project>

.idea/model_scan.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/model_scan.iml" filepath="$PROJECT_DIR$/.idea/model_scan.iml" />
+    </modules>
+  </component>
+</project>

Dockerfile

@@ -0,0 +1,15 @@
+FROM clamav/clamav:stable
+
+RUN apk update && apk add --no-cache \
+    python3 \
+    py3-pip \
+    && rm -rf /var/cache/apk/*
+
+WORKDIR /app
+
+COPY requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt
+
+COPY . .
+
+CMD ["sh", "-c", "freshclam && python3 scan_main.py"]

model.py

@@ -0,0 +1,48 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class ImageScanRequest(BaseModel):
+    imageId: int
+    url: str
+    wait: bool
+    scans: List[int]
+    callbackUrl: str
+
+
+class ImageScanTag(BaseModel):
+    type: str
+    name: str
+
+
+class ImageScanResponse(BaseModel):
+    ok: bool
+    error: str
+    deleted: bool
+    blockedFor: List[str]
+    tags: List[ImageScanTag]
+
+
+class ImageTag(BaseModel):
+    tag: str
+    id: Optional[int]
+    confidence: int
+
+
+class ImageScanCallbackRequest(BaseModel):
+    id: int
+    isValid: bool
+    tags: List[ImageTag]
+
+
+class ModelScanRequest(BaseModel):
+    callbackUrl: str
+    fileUrl: str
+    lowPriority: bool
+    tasks: List[str]
+
+
+class ModelScanResponse(BaseModel):
+    ok: bool
+    error: str

pure_blake3.py

@@ -0,0 +1,353 @@
+#! /usr/bin/env python3
+
+# This is a Python port of the Rust reference implementation of BLAKE3:
+# https://github.com/BLAKE3-team/BLAKE3/blob/master/reference_impl/reference_impl.rs
+
+from __future__ import annotations
+from dataclasses import dataclass
+
+OUT_LEN = 32
+KEY_LEN = 32
+BLOCK_LEN = 64
+CHUNK_LEN = 1024
+
+CHUNK_START = 1 << 0
+CHUNK_END = 1 << 1
+PARENT = 1 << 2
+ROOT = 1 << 3
+KEYED_HASH = 1 << 4
+DERIVE_KEY_CONTEXT = 1 << 5
+DERIVE_KEY_MATERIAL = 1 << 6
+
+IV = [
+    0x6A09E667,
+    0xBB67AE85,
+    0x3C6EF372,
+    0xA54FF53A,
+    0x510E527F,
+    0x9B05688C,
+    0x1F83D9AB,
+    0x5BE0CD19,
+]
+
+MSG_PERMUTATION = [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8]
+
+
+def mask32(x: int) -> int:
+    return x & 0xFFFFFFFF
+
+
+def add32(x: int, y: int) -> int:
+    return mask32(x + y)
+
+
+def rightrotate32(x: int, n: int) -> int:
+    return mask32(x << (32 - n)) | (x >> n)
+
+
+# The mixing function, G, which mixes either a column or a diagonal.
+def g(state: list[int], a: int, b: int, c: int, d: int, mx: int, my: int) -> None:
+    state[a] = add32(state[a], add32(state[b], mx))
+    state[d] = rightrotate32(state[d] ^ state[a], 16)
+    state[c] = add32(state[c], state[d])
+    state[b] = rightrotate32(state[b] ^ state[c], 12)
+    state[a] = add32(state[a], add32(state[b], my))
+    state[d] = rightrotate32(state[d] ^ state[a], 8)
+    state[c] = add32(state[c], state[d])
+    state[b] = rightrotate32(state[b] ^ state[c], 7)
+
+
+def round(state: list[int], m: list[int]) -> None:
+    # Mix the columns.
+    g(state, 0, 4, 8, 12, m[0], m[1])
+    g(state, 1, 5, 9, 13, m[2], m[3])
+    g(state, 2, 6, 10, 14, m[4], m[5])
+    g(state, 3, 7, 11, 15, m[6], m[7])
+    # Mix the diagonals.
+    g(state, 0, 5, 10, 15, m[8], m[9])
+    g(state, 1, 6, 11, 12, m[10], m[11])
+    g(state, 2, 7, 8, 13, m[12], m[13])
+    g(state, 3, 4, 9, 14, m[14], m[15])
+
+
+def permute(m: list[int]) -> None:
+    original = list(m)
+    for i in range(16):
+        m[i] = original[MSG_PERMUTATION[i]]
+
+
+def compress(
+    chaining_value: list[int],
+    block_words: list[int],
+    counter: int,
+    block_len: int,
+    flags: int,
+) -> list[int]:
+    state = [
+        chaining_value[0],
+        chaining_value[1],
+        chaining_value[2],
+        chaining_value[3],
+        chaining_value[4],
+        chaining_value[5],
+        chaining_value[6],
+        chaining_value[7],
+        IV[0],
+        IV[1],
+        IV[2],
+        IV[3],
+        mask32(counter),
+        mask32(counter >> 32),
+        block_len,
+        flags,
+    ]
+
+    assert len(block_words) == 16
+    block = list(block_words)
+
+    round(state, block)  # round 1
+    permute(block)
+    round(state, block)  # round 2
+    permute(block)
+    round(state, block)  # round 3
+    permute(block)
+    round(state, block)  # round 4
+    permute(block)
+    round(state, block)  # round 5
+    permute(block)
+    round(state, block)  # round 6
+    permute(block)
+    round(state, block)  # round 7
+
+    for i in range(8):
+        state[i] ^= state[i + 8]
+        state[i + 8] ^= chaining_value[i]
+
+    return state
+
+
+def words_from_little_endian_bytes(b: bytes) -> list[int]:
+    assert len(b) % 4 == 0
+    return [int.from_bytes(b[i : i + 4], "little") for i in range(0, len(b), 4)]
+
+
+# Each chunk or parent node can produce either an 8-word chaining value or, by
+# setting the ROOT flag, any number of final output bytes. The Output struct
+# captures the state just prior to choosing between those two possibilities.
+@dataclass
+class Output:
+    input_chaining_value: list[int]
+    block_words: list[int]
+    counter: int
+    block_len: int
+    flags: int
+
+    def chaining_value(self) -> list[int]:
+        return compress(
+            self.input_chaining_value,
+            self.block_words,
+            self.counter,
+            self.block_len,
+            self.flags,
+        )[:8]
+
+    def root_output_bytes(self, length: int) -> bytes:
+        output_bytes = bytearray()
+        i = 0
+        while i < length:
+            words = compress(
+                self.input_chaining_value,
+                self.block_words,
+                i // 64,
+                self.block_len,
+                self.flags | ROOT,
+            )
+            # The output length might not be a multiple of 4.
+            for word in words:
+                word_bytes = word.to_bytes(4, "little")
+                take = min(len(word_bytes), length - i)
+                output_bytes.extend(word_bytes[:take])
+                i += take
+        return output_bytes
+
+
+@dataclass
+class ChunkState:
+    chaining_value: list[int]
+    chunk_counter: int
+    block: bytearray
+    block_len: int
+    blocks_compressed: int
+    flags: int
+
+    def __init__(self, key_words: list[int], chunk_counter: int, flags: int) -> None:
+        self.chaining_value = key_words
+        self.chunk_counter = chunk_counter
+        self.block = bytearray(BLOCK_LEN)
+        self.block_len = 0
+        self.blocks_compressed = 0
+        self.flags = flags
+
+    def len(self) -> int:
+        return BLOCK_LEN * self.blocks_compressed + self.block_len
+
+    def start_flag(self) -> int:
+        if self.blocks_compressed == 0:
+            return CHUNK_START
+        else:
+            return 0
+
+    def update(self, input_bytes: bytes) -> None:
+        while input_bytes:
+            # If the block buffer is full, compress it and clear it. More
+            # input_bytes is coming, so this compression is not CHUNK_END.
+            if self.block_len == BLOCK_LEN:
+                block_words = words_from_little_endian_bytes(self.block)
+                self.chaining_value = compress(
+                    self.chaining_value,
+                    block_words,
+                    self.chunk_counter,
+                    BLOCK_LEN,
+                    self.flags | self.start_flag(),
+                )[:8]
+                self.blocks_compressed += 1
+                self.block = bytearray(BLOCK_LEN)
+                self.block_len = 0
+
+            # Copy input bytes into the block buffer.
+            want = BLOCK_LEN - self.block_len
+            take = min(want, len(input_bytes))
+            self.block[self.block_len : self.block_len + take] = input_bytes[:take]
+            self.block_len += take
+            input_bytes = input_bytes[take:]
+
+    def output(self) -> Output:
+        block_words = words_from_little_endian_bytes(self.block)
+        return Output(
+            self.chaining_value,
+            block_words,
+            self.chunk_counter,
+            self.block_len,
+            self.flags | self.start_flag() | CHUNK_END,
+        )
+
+
+def parent_output(
+    left_child_cv: list[int],
+    right_child_cv: list[int],
+    key_words: list[int],
+    flags: int,
+) -> Output:
+    return Output(
+        key_words, left_child_cv + right_child_cv, 0, BLOCK_LEN, PARENT | flags
+    )
+
+
+def parent_cv(
+    left_child_cv: list[int],
+    right_child_cv: list[int],
+    key_words: list[int],
+    flags: int,
+) -> list[int]:
+    return parent_output(
+        left_child_cv, right_child_cv, key_words, flags
+    ).chaining_value()
+
+
+# An incremental hasher that can accept any number of writes.
+@dataclass
+class Hasher:
+    chunk_state: ChunkState
+    key_words: list[int]
+    cv_stack: list[list[int]]
+    flags: int
+
+    def _init(self, key_words: list[int], flags: int) -> None:
+        assert len(key_words) == 8
+        self.chunk_state = ChunkState(key_words, 0, flags)
+        self.key_words = key_words
+        self.cv_stack = []
+        self.flags = flags
+
+    # Construct a new `Hasher` for the regular hash function.
+    def __init__(self) -> None:
+        self._init(IV, 0)
+
+    # Construct a new `Hasher` for the keyed hash function.
+    @classmethod
+    def new_keyed(cls, key: bytes) -> Hasher:
+        keyed_hasher = cls()
+        key_words = words_from_little_endian_bytes(key)
+        keyed_hasher._init(key_words, KEYED_HASH)
+        return keyed_hasher
+
+    # Construct a new `Hasher` for the key derivation function. The context
+    # string should be hardcoded, globally unique, and application-specific.
+    @classmethod
+    def new_derive_key(cls, context: str) -> Hasher:
+        context_hasher = cls()
+        context_hasher._init(IV, DERIVE_KEY_CONTEXT)
+        context_hasher.update(context.encode("utf8"))
+        context_key = context_hasher.finalize(KEY_LEN)
+        context_key_words = words_from_little_endian_bytes(context_key)
+        derive_key_hasher = cls()
+        derive_key_hasher._init(context_key_words, DERIVE_KEY_MATERIAL)
+        return derive_key_hasher
+
+    # Section 5.1.2 of the BLAKE3 spec explains this algorithm in more detail.
+    def add_chunk_chaining_value(self, new_cv: list[int], total_chunks: int) -> None:
+        # This chunk might complete some subtrees. For each completed subtree,
+        # its left child will be the current top entry in the CV stack, and
+        # its right child will be the current value of `new_cv`. Pop each left
+        # child off the stack, merge it with `new_cv`, and overwrite `new_cv`
+        # with the result. After all these merges, push the final value of
+        # `new_cv` onto the stack. The number of completed subtrees is given
+        # by the number of trailing 0-bits in the new total number of chunks.
+        while total_chunks & 1 == 0:
+            new_cv = parent_cv(self.cv_stack.pop(), new_cv, self.key_words, self.flags)
+            total_chunks >>= 1
+        self.cv_stack.append(new_cv)
+
+    # Add input to the hash state. This can be called any number of times.
+    def update(self, input_bytes: bytes) -> None:
+        while input_bytes:
+            # If the current chunk is complete, finalize it and reset the
+            # chunk state. More input is coming, so this chunk is not ROOT.
+            if self.chunk_state.len() == CHUNK_LEN:
+                chunk_cv = self.chunk_state.output().chaining_value()
+                total_chunks = self.chunk_state.chunk_counter + 1
+                self.add_chunk_chaining_value(chunk_cv, total_chunks)
+                self.chunk_state = ChunkState(self.key_words, total_chunks, self.flags)
+
+            # Compress input bytes into the current chunk state.
+            want = CHUNK_LEN - self.chunk_state.len()
+            take = min(want, len(input_bytes))
+            self.chunk_state.update(input_bytes[:take])
+            input_bytes = input_bytes[take:]
+
+    # Finalize the hash and write any number of output bytes.
+    def finalize(self, length: int = OUT_LEN) -> bytes:
+        # Starting with the Output from the current chunk, compute all the
+        # parent chaining values along the right edge of the tree, until we
+        # have the root Output.
+        output = self.chunk_state.output()
+        parent_nodes_remaining = len(self.cv_stack)
+        while parent_nodes_remaining > 0:
+            parent_nodes_remaining -= 1
+            output = parent_output(
+                self.cv_stack[parent_nodes_remaining],
+                output.chaining_value(),
+                self.key_words,
+                self.flags,
+            )
+        return output.root_output_bytes(length)
+
+
+# If this file is executed directly, hash standard input.
+if __name__ == "__main__":
+    import sys
+
+    hasher = Hasher()
+    while buf := sys.stdin.buffer.read(65536):
+        hasher.update(buf)
+    print(hasher.finalize().hex())

requirements.txt

@@ -0,0 +1,6 @@
+crcmod
+pyclamd
+uvicorn[standard]
+fastapi
+requests
+picklescan

scan.py

@@ -0,0 +1,76 @@
+import pyclamd
+
+from picklescan.scanner import (
+    scan_url,
+    ScanResult, SafetyLevel
+)
+
+
+# def scan_file(file_path: str):
+#     ret = scan_pickle_bytes(io.BytesIO(pickle.dumps(file_path)), "file.pkl")
+#     print(ret)
+
+
+def scan_file(file_path: str):
+    if file_path.startswith("http"):
+        scan_result: ScanResult = scan_url(file_path)
+    else:
+        return None
+
+    globalImports = list(map(lambda x: fmt_import(x.module, x.name), scan_result.globals))
+    dangerousImports = list(map(lambda x: fmt_import(x.module, x.name),
+                                filter(lambda x: x.safety == SafetyLevel.Dangerous, scan_result.globals)))
+    if len(dangerousImports) > 0:
+        picklescanExitCode = 1
+    else:
+        picklescanExitCode = 0
+    return {
+        'url': file_path,
+        'fileExists': True,
+        'picklescanExitCode': picklescanExitCode,
+        'picklescanGlobalImports': globalImports,
+        'picklescanDangerousImports': dangerousImports,
+        # 'clamscanExitCode': ScanExitCode,
+        # 'clamscanOutput': string,
+        # hashes: Record < ModelHashType, string >;
+        # conversions: Record < 'safetensors' | 'ckpt', ConversionResult >;
+    }
+
+
+def init_clamd():
+    clamd = pyclamd.ClamdUnixSocket()
+    return clamd
+
+
+def clamd_file(file_path: str, clamd):
+    if file_path.startswith("http"):
+        import urllib.request
+        tmp_path = f'/tmp/clamd_{file_path.split("/")[-1]}'
+        urllib.request.urlretrieve(file_path, tmp_path)
+        ret = clamd.scan_file(tmp_path)
+        if ret is None:
+            return {
+                'clamscanExitCode': 0,
+                'clamscanOutput': "No virus found",
+            }
+        elif file_path in ret and len(file_path) > 0:
+            return {
+                'clamscanExitCode': 1,
+                'clamscanOutput': ' '.join(ret[file_path]),
+            }
+
+
+def fmt_import(module: str, name: str):
+    return f"from ${module} import ${name}",
+
+
+if __name__ == "__main__":
+    detail = scan_file("https://huggingface.co/yesyeahvh/bad-hands-5/resolve/main/bad-hands-5.pt")
+    clamd_detail = clamd_file("https://huggingface.co/yesyeahvh/bad-hands-5/resolve/main/bad-hands-5.pt")
+    print(detail)
+    print(clamd_detail)
+    # ScanResult(
+    #     globals=[Global(module='torch', name='FloatStorage', safety= < SafetyLevel.Innocuous: 'innocuous' >),
+    #              Global(module='collections', name='OrderedDict', safety= < SafetyLevel.Innocuous: 'innocuous' >),
+    #              Global(module='torch._utils', name='_rebuild_tensor_v2',safety= < SafetyLevel.Innocuous: 'innocuous' >)],
+    #     scanned_files = 1, issues_count = 0, infected_files = 0, scan_err = False)

scan_convert.py

@@ -0,0 +1,4 @@
+
+
+def convert_and_upload(from_type: str, target_type: str):
+    pass

scan_hash.py

@@ -0,0 +1,68 @@
+import hashlib
+import os.path
+import sys
+from typing import Dict
+import binascii
+import crcmod
+import pure_blake3
+
+
+def get_hash_string(hash_bytes: bytes) -> str:
+    return binascii.hexlify(hash_bytes).decode()
+
+
+def compute_AutoV1Hash(file_stream) -> str:
+    minFileSize = 0x100000 * 2
+    if file_stream.seek(0, 2) < minFileSize:
+        return None
+    file_stream.seek(0x100000)
+    buffer = file_stream.read(0x10000)
+    hashBytes = hashlib.sha256(buffer).digest()
+    hashString = get_hash_string(hashBytes)
+    return hashString[:8]
+
+
+def ComputeCRC32Hash(file_stream) -> str:
+    crc32 = crcmod.predefined.Crc('crc-32c')
+    file_stream.seek(0)
+    for chunk in iter(lambda: file_stream.read(4096), b""):
+        crc32.update(chunk)
+    return get_hash_string(crc32.digest())
+
+
+def generate_model_hashes(file_path: str) -> Dict[str, str]:
+    if file_path.startswith("http"):
+        import urllib.request
+        tmp_path = f'/tmp/clamd_{file_path.split("/")[-1]}'
+        if os.path.exists(tmp_path):
+            pass
+        else:
+            urllib.request.urlretrieve(file_path, tmp_path)
+        file_path = tmp_path
+
+    sha256 = hashlib.sha256()
+    blake3Hasher = pure_blake3.Hasher()
+
+    with open(file_path, "rb") as fileStream:
+        for chunk in iter(lambda: fileStream.read(4096), b""):
+            sha256.update(chunk)
+            blake3Hasher.update(chunk)
+    sha256HashString = get_hash_string(sha256.digest())
+    autoV1HashString = compute_AutoV1Hash(open(file_path, "rb"))
+    autoV2HashString = sha256HashString[:10]
+    blake3HashString = blake3Hasher.finalize().hex()
+    crc32HashString = ComputeCRC32Hash(open(file_path, "rb"))
+
+    result = {
+        "SHA256": sha256HashString,
+        "AutoV1": autoV1HashString,
+        "AutoV2": autoV2HashString,
+        "BLAKE3": blake3HashString,
+        "CRC32": crc32HashString,
+    }
+
+    return result
+
+
+if __name__ == "__main__":
+    print(generate_model_hashes(".gitignore"))

scan_main.py

@@ -0,0 +1,36 @@
+import requests
+import uvicorn
+from fastapi import BackgroundTasks, FastAPI
+import model
+from scan import init_clamd, scan_file, clamd_file
+from scan_hash import generate_model_hashes
+
+app = FastAPI()
+
+
+def write_scan_model_result(req: model.ModelScanRequest):
+    ret = {}
+    if 'Scan' in req.tasks:
+        detail = scan_file(req.fileUrl)
+        clamd_detail = clamd_file(req.fileUrl, clamd_exec)
+        ret |= detail
+        ret |= clamd_detail
+    if 'Hash' in req.tasks:
+        ret.hashes = generate_model_hashes(req.fileUrl)
+
+    try:
+        requests.post(req.callbackUrl, json=ret)
+    except Exception as ex:
+        print(ex)
+
+
+@app.post("/model-scan")
+async def model_scan_handler(req: model.ModelScanRequest, background_tasks: BackgroundTasks):
+    background_tasks.add_task(write_scan_model_result, req)
+    return model.ModelScanResponse(ok=True, error="")
+
+
+if __name__ == "__main__":
+    global clamd_exec
+    clamd_exec = init_clamd()
+    uvicorn.run(app, host="0.0.0.0", port=7860)

test_main.http

@@ -0,0 +1,11 @@
+# Test your FastAPI endpoints
+
+GET http://127.0.0.1:8000/
+Accept: application/json
+
+###
+
+GET http://127.0.0.1:8000/hello/User
+Accept: application/json
+
+###