310 lines
9.7 KiB
Python
310 lines
9.7 KiB
Python
from PYAS_Diagnostics import log_exception
|
|
import os
|
|
import time
|
|
import threading
|
|
import requests
|
|
|
|
|
|
class CloudScanner:
|
|
def __init__(self):
|
|
self.api_host = None
|
|
self.api_key = None
|
|
self.timeout = 30
|
|
|
|
self.lock = threading.RLock()
|
|
self.local = threading.local()
|
|
|
|
def _get_session(self, api_host, api_key):
|
|
if (
|
|
not hasattr(self.local, "session")
|
|
or getattr(self.local, "host", None) != api_host
|
|
or getattr(self.local, "key", None) != api_key
|
|
):
|
|
previous = getattr(self.local, "session", None)
|
|
|
|
if previous is not None:
|
|
previous.close()
|
|
|
|
session = requests.Session()
|
|
session.headers.update({"X-API-Key": api_key, "User-Agent": "PYAS-Engine/1.1"})
|
|
self.local.session = session
|
|
self.local.host = api_host
|
|
self.local.key = api_key
|
|
|
|
return self.local.session
|
|
|
|
def _request(self, method, endpoint, api_host, api_key, **kwargs):
|
|
session = self._get_session(api_host, api_key)
|
|
|
|
try:
|
|
r = session.request(method, f"{api_host}{endpoint}", timeout=self.timeout, **kwargs)
|
|
|
|
if r.status_code == 200:
|
|
return r
|
|
|
|
except Exception:
|
|
log_exception("PYAS_Cloud.CloudScanner._request:31")
|
|
pass
|
|
|
|
return None
|
|
|
|
def rescan(self, sha256, api_host, api_key):
|
|
r = self._request("POST", f"/api/rescan/{sha256}", api_host, api_key)
|
|
return r is not None and r.json().get("status") == "success"
|
|
|
|
def upload_file(
|
|
self,
|
|
file_path,
|
|
api_host,
|
|
api_key,
|
|
chunk_size=4194304,
|
|
need_rescan=False,
|
|
max_retries=3,
|
|
file_hash=None,
|
|
cancel_event=None,
|
|
):
|
|
try:
|
|
sha256 = file_hash
|
|
cancel_event = cancel_event or threading.Event()
|
|
|
|
if cancel_event.is_set() or not sha256:
|
|
return False, None
|
|
|
|
status_req = self._request("GET", f"/api/processing_status/{sha256}", api_host, api_key)
|
|
|
|
if cancel_event.is_set():
|
|
return False, sha256
|
|
|
|
if status_req:
|
|
current_status = status_req.json().get("status")
|
|
|
|
if current_status == "done":
|
|
if need_rescan:
|
|
self.rescan(sha256, api_host, api_key)
|
|
|
|
return True, sha256
|
|
|
|
elif current_status in ("queued", "processing"):
|
|
return True, sha256
|
|
|
|
file_size = os.path.getsize(file_path)
|
|
|
|
if file_size > 104857600:
|
|
return False, sha256
|
|
|
|
total_chunks = max(1, (file_size + chunk_size - 1) // chunk_size)
|
|
upload_id = os.urandom(16).hex()
|
|
|
|
with open(file_path, "rb") as f:
|
|
for i in range(total_chunks):
|
|
if cancel_event.is_set():
|
|
return False, sha256
|
|
|
|
chunk_data = f.read(chunk_size)
|
|
headers = {
|
|
"X-Chunk-Index": str(i),
|
|
"X-Total-Chunks": str(total_chunks),
|
|
"X-Upload-ID": upload_id,
|
|
}
|
|
|
|
chunk_success = False
|
|
|
|
for attempt in range(max_retries):
|
|
if cancel_event.is_set():
|
|
return False, sha256
|
|
|
|
r = self._request(
|
|
"POST",
|
|
"/api/upload",
|
|
api_host,
|
|
api_key,
|
|
files={"file": (os.path.basename(file_path), chunk_data)},
|
|
headers=headers,
|
|
)
|
|
|
|
if cancel_event.is_set():
|
|
return False, sha256
|
|
|
|
if r:
|
|
if i == total_chunks - 1:
|
|
try:
|
|
resp = r.json()
|
|
|
|
if "url" in resp:
|
|
sha256 = resp.get("url", "").split("/")[-1]
|
|
except Exception:
|
|
log_exception("PYAS_Cloud.CloudScanner.upload_file:89")
|
|
pass
|
|
|
|
chunk_success = True
|
|
break
|
|
|
|
if cancel_event.wait(2**attempt):
|
|
return False, sha256
|
|
|
|
if not chunk_success:
|
|
return False, sha256
|
|
|
|
return True, sha256
|
|
except Exception:
|
|
log_exception("PYAS_Cloud.CloudScanner.upload_file:101")
|
|
return False, None
|
|
|
|
def get_result(self, sha256, api_host, api_key, max_retries=6, interval=10):
|
|
try:
|
|
if not sha256:
|
|
return False
|
|
|
|
is_done = False
|
|
|
|
for _ in range(max_retries):
|
|
r = self._request("GET", f"/api/processing_status/{sha256}", api_host, api_key)
|
|
|
|
if r:
|
|
st = r.json().get("status", "error")
|
|
|
|
if st == "done":
|
|
is_done = True
|
|
break
|
|
|
|
if st in ["error", "failed"]:
|
|
return False
|
|
|
|
time.sleep(interval)
|
|
|
|
if not is_done:
|
|
return False
|
|
|
|
r = self._request("GET", f"/api/report/{sha256}", api_host, api_key)
|
|
|
|
if r:
|
|
data = r.json().get("data", {})
|
|
metadata = data.get("metadata", {})
|
|
label = metadata.get("label", "Unsupport")
|
|
score = metadata.get("score", 0)
|
|
sims = data.get("similar", [])
|
|
|
|
is_malicious = "General" in label
|
|
sim_malicious_count = 0
|
|
valid_sim_count = 0
|
|
|
|
for s in sims:
|
|
if s.get("similarity", 0) > 80:
|
|
valid_sim_count += 1
|
|
|
|
if "General" in s.get("label", ""):
|
|
sim_malicious_count += 1
|
|
|
|
if is_malicious and (
|
|
valid_sim_count == 0 or sim_malicious_count == valid_sim_count
|
|
):
|
|
return f"Malware:WinPE/General.{score}!cl"
|
|
|
|
except Exception:
|
|
log_exception("PYAS_Cloud.CloudScanner.get_result:149")
|
|
pass
|
|
|
|
return False
|
|
|
|
|
|
class CloudQueueMixin:
|
|
def cloud_check(self, file_path):
|
|
with self.lock_config:
|
|
if not self.pyas_config.get("cloud_switch", False):
|
|
return
|
|
|
|
cancel_event = self.cloud_cancel_event
|
|
|
|
norm_path = self.norm_path(file_path)
|
|
|
|
if not norm_path or cancel_event.is_set():
|
|
return
|
|
|
|
cache_key = (os.path.normcase(norm_path), cancel_event)
|
|
|
|
with self.lock_file_ops:
|
|
if cache_key in self.cloud_pending:
|
|
return
|
|
|
|
self.cloud_pending.add(cache_key)
|
|
|
|
self.cloud_queue.put((norm_path, cancel_event))
|
|
|
|
def cloud_worker(self):
|
|
while True:
|
|
try:
|
|
file_path, cancel_event = self.cloud_queue.get()
|
|
cache_key = (os.path.normcase(file_path), cancel_event)
|
|
|
|
try:
|
|
self.perform_cloud_scan(file_path, cancel_event)
|
|
finally:
|
|
with self.lock_file_ops:
|
|
self.cloud_pending.discard(cache_key)
|
|
|
|
self.cloud_queue.task_done()
|
|
except Exception:
|
|
log_exception("PYAS_Cloud.CloudQueueMixin.cloud_worker:31")
|
|
pass
|
|
|
|
def perform_cloud_scan(self, file_path, cancel_event=None):
|
|
was_locked = False
|
|
|
|
try:
|
|
with self.lock_config:
|
|
if not self.pyas_config.get("cloud_switch", False):
|
|
return False
|
|
|
|
cancel_event = cancel_event or self.cloud_cancel_event
|
|
api_host, api_key, max_size = (
|
|
self.pyas_config.get("api_host"),
|
|
self.pyas_config.get("api_key"),
|
|
self.pyas_config.get("size", 256 * 1024 * 1024),
|
|
)
|
|
|
|
if (
|
|
cancel_event.is_set()
|
|
or not os.path.exists(file_path)
|
|
or not os.path.isfile(file_path)
|
|
):
|
|
return False
|
|
|
|
with self.lock_file_ops:
|
|
if file_path in self.virus_lock:
|
|
self.lock_file(file_path, False)
|
|
was_locked = True
|
|
|
|
if os.path.getsize(file_path) > max_size:
|
|
if was_locked:
|
|
self.lock_file(file_path, True)
|
|
|
|
return False
|
|
|
|
file_hash = self.calc_file_hash(file_path)
|
|
success, sha256 = self.cloud.upload_file(
|
|
file_path, api_host, api_key, file_hash=file_hash, cancel_event=cancel_event
|
|
)
|
|
|
|
if was_locked:
|
|
self.lock_file(file_path, True)
|
|
was_locked = False
|
|
|
|
if not success and not cancel_event.is_set():
|
|
self.write_log(
|
|
"WARN", "Cloud API", source=file_path, detail="Failed", success=False
|
|
)
|
|
|
|
except Exception as e:
|
|
log_exception("PYAS_Cloud.CloudQueueMixin.perform_cloud_scan:66")
|
|
self.write_log("WARN", "perform_cloud_scan", detail=str(e), success=False)
|
|
|
|
finally:
|
|
if was_locked:
|
|
try:
|
|
self.lock_file(file_path, True)
|
|
except Exception:
|
|
log_exception("PYAS_Cloud.CloudQueueMixin.perform_cloud_scan:73")
|
|
pass
|
|
|
|
return False
|