pengdaqian commited on
Commit
301572e
1 Parent(s): 084eefe
Files changed (1) hide show
  1. scan_hash.py +10 -2
scan_hash.py CHANGED
@@ -4,6 +4,8 @@ import sys
4
  from typing import Dict
5
  import binascii
6
  import crcmod
 
 
7
  import pure_blake3
8
 
9
 
@@ -30,14 +32,20 @@ def ComputeCRC32Hash(file_stream) -> str:
30
  return get_hash_string(crc32.digest())
31
 
32
 
 
 
 
 
33
  def generate_model_hashes(file_path: str) -> Dict[str, str]:
34
  if file_path.startswith("http"):
35
  import urllib.request
36
- tmp_path = f'/tmp/clamd_{file_path.split("/")[-1]}'
37
  if os.path.exists(tmp_path):
38
  pass
39
  else:
40
- urllib.request.urlretrieve(file_path, tmp_path)
 
 
41
  file_path = tmp_path
42
 
43
  sha256 = hashlib.sha256()
 
4
  from typing import Dict
5
  import binascii
6
  import crcmod
7
+ import requests
8
+
9
  import pure_blake3
10
 
11
 
 
32
  return get_hash_string(crc32.digest())
33
 
34
 
35
+ headers = {
36
+ "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/78.0.3904.108 Safari/537.36"
37
+ }
38
+
39
  def generate_model_hashes(file_path: str) -> Dict[str, str]:
40
  if file_path.startswith("http"):
41
  import urllib.request
42
+ tmp_path = f'/tmp/clamd_{file_path.split("/")[-1].split("?")[0]}'
43
  if os.path.exists(tmp_path):
44
  pass
45
  else:
46
+ resp = requests.get(file_path, headers=headers).content
47
+ with open(tmp_path, "wb") as f:
48
+ f.write(resp)
49
  file_path = tmp_path
50
 
51
  sha256 = hashlib.sha256()