Add files via upload

PYAS.py

@@ -19,7 +19,7 @@ class MainWindow_Controller(QMainWindow):
         self.setWindowFlags(Qt.FramelessWindowHint)
         self.pyas = sys.argv[0].replace("\\", "/")
         self.dir = os.path.dirname(self.pyas)
-        self.pyae_version = "PE Simhash Engine"
+        self.pyae_version = "PE SimHash Engine"
         self.pyas_version = "3.0.4"
         self.ui = Ui_MainWindow()
         self.ui.setupUi(self)
@@ -627,8 +627,8 @@ class MainWindow_Controller(QMainWindow):
     def update_database(self):
         try:
             if self.question_event("您確定要更新數據庫嗎?"):
-                file_path = os.path.join(self.dir, "PYAS_Function.vdb")
-                response = requests.get('http://27.147.30.238:5001/database', params={"version": self.pyae_version}, timeout=5)
+                file_path = os.path.join(self.dir, "PYAS_Model.json")
+                response = requests.get('http://27.147.30.238:5001/model', params={"version": self.pyae_version}, timeout=10)
                 if response.status_code == 200:
                     with open(file_path, 'wb') as f:
                         f.write(response.content)
@@ -906,11 +906,10 @@ class MainWindow_Controller(QMainWindow):
                             pass
             QApplication.processEvents()
             ones, zero = self.pe.predict(fn)
-            #print(ones, zero)
             if self.json["high_sensitive"]:
-                return ones >= zero
+                return ones == 1.0 or ones >= zero
             elif "_CorExeMain" not in fn:
-                return ones - zero > 0.1
+                return ones == 1.0 or ones - zero >= 0.1
             return False
         except:
             return False

PYAS_Engine.py

@@ -1,4 +1,4 @@
-import json, re, hashlib, numpy
+import json, re, hashlib, numpy, time
 from itertools import groupby
 
 class ListSimHash:
@@ -16,9 +16,14 @@ class ListSimHash:
                 model_data = json.load(f)
         self.model = model_data
 
-    def train_model(self, ones, zero, feature=128, batch_size=200):
+    def get_model(self, label):
+        return self.model[label]
+
+    def train_model(self, ones, zero, feature=128, batch_size=1000, progress=True):
+        self.progress = progress
         self.model["feature"] = feature
         self.model["batch_size"] = batch_size
+        self.model["version"] = time.strftime("%Y-%m-%d %H:%M:%S")
         self.model["ones"] = [self.build_text(" ".join(x)) for x in ones]
         self.model["zero"] = [self.build_text(" ".join(y)) for y in zero]
 
@@ -27,15 +32,16 @@ class ListSimHash:
         content = ''.join(re.findall(r'[\w\u4e00-\u9fcc]+', content.lower()))
         features = [content[i:i + width] for i in range(max(len(content) - width + 1, 1))]
         features = {k: sum(1 for _ in g) for k, g in groupby(sorted(features))}
-        for f, w in features.items():
+        total = len(features)
+        for i, (f, w) in enumerate(features.items(), 1):
             count += w
             h = hashlib.md5(f.encode('utf-8')).digest()
             batch.append(h * w)
-            if len(batch) >= self.model["batch_size"]:
+            if len(batch) >= self.model["batch_size"] or i == total:
                 sums.append(self.sum_hashes(batch))
                 batch = []
-        if batch:
-            sums.append(self.sum_hashes(batch))
+                if self.progress:
+                    self.progress_bar(i, total, prefix=f'Block {i}/{total}:')
         combined_sums = numpy.sum(sums, 0)
         v = numpy.packbits(combined_sums > count / 2).tobytes()
         return int.from_bytes(v, 'big')
@@ -44,7 +50,16 @@ class ListSimHash:
         bitarray = numpy.unpackbits(numpy.frombuffer(b''.join(digests), dtype='>B'))
         return numpy.sum(numpy.reshape(bitarray, (-1, self.model["feature"])), 0)
 
-    def predict(self, query):
+    def progress_bar(self, iteration, total, prefix='', suffix='', length=50, fill='█'):
+        percent = min(100.0, max(0.0, 100 * (iteration / float(total))))
+        end_char = '\n' if percent >= 100 else '\r'
+        percent_string = "{0:.2f}".format(percent)
+        filled_length = int(length * iteration // total)
+        bar = fill * filled_length + ' ' * (length - filled_length)
+        print(f'\r{prefix: <20} |{bar}| {percent_string}% {suffix}', end=end_char)
+
+    def predict(self, query, progress=False):
+        self.progress = progress
         query_hash = self.build_text(" ".join(query))
         max_ones = max(self.similar(x, query_hash) for x in self.model["ones"])
         max_zero = max(self.similar(y, query_hash) for y in self.model["zero"])