diff --git a/PYAS.py b/PYAS.py index 9991574..74f189d 100644 --- a/PYAS.py +++ b/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 diff --git a/PYAS_Engine.py b/PYAS_Engine.py index c48b629..c616492 100644 --- a/PYAS_Engine.py +++ b/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"])