mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-08-22 08:43:54 +08:00
firestar5683
39 lines
1.7 KiB
Python
39 lines
1.7 KiB
Python
from tinygrad.tensor import Tensor
|
|
from tinygrad.helpers import CHUNK_SIZE
|
|
from tinygrad.nn.state import fs_load
|
|
import argparse, math, hashlib
|
|
|
|
def _python_hash_1mb(data:bytes|bytearray):
|
|
chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
|
|
chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
|
|
return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
|
|
|
|
def hash_file(data: bytes|bytearray):
|
|
if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
|
|
base_chunks = math.ceil(len(data) / CHUNK_SIZE)
|
|
tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
|
|
|
|
for _ in range(tree_depth + 1):
|
|
data_chunks = [data[i:i+CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
|
|
data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
|
|
data = b''.join(data_chunk_hashes)
|
|
if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
|
|
|
|
return data[:16]
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
|
|
parser.add_argument("--len", type=int, required=True, help="file length to fetch")
|
|
parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
|
|
parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
|
|
args = parser.parse_args()
|
|
|
|
fs_load(Tensor(bytes.fromhex(args.hash), device="CPU"), args.len).to(f"disk:{args.dest}").realize()
|
|
|
|
if args.check:
|
|
with open(args.dest, "rb") as f:
|
|
data = f.read()
|
|
assert hash_file(data) == bytes.fromhex(args.hash), "Hash mismatch after fetching file"
|
|
print("File hash verified successfully!")
|