nirizr · Feb 27, 2017
diff --git a/‎idaplugin/rematch/actions/match.py
+135-48 b/‎idaplugin/rematch/actions/match.py
+135-48
diff --git a/‎idaplugin/rematch/dialogs/__init__.py
+2-1 b/‎idaplugin/rematch/dialogs/__init__.py
+2-1
diff --git a/‎idaplugin/rematch/dialogs/base.py
+5-5 b/‎idaplugin/rematch/dialogs/base.py
+5-5
@@ -2,6 +2,7 @@
 import idautils
 
 from ..dialogs.match import MatchDialog
+from ..dialogs.matchresult import MatchResultDialog
 
 from .. import instances
 from .. import network, netnode, log
@@ -17,8 +18,7 @@ class MatchAction(base.BoundFileAction):
   def __init__(self, *args, **kwargs):
     super(MatchAction, self).__init__(*args, **kwargs)
     self.functions = None
-    self.pbar = None
-    self.timer = None
+    self.results = None
     self.task_id = None
     self.file_version_id = None
     self.instance_set = []
@@ -31,6 +31,37 @@ def __init__(self, *args, **kwargs):
     self.target_file = None
     self.methods = None
 
+    self.delayed_queries = []
+
+    self.pbar = QtWidgets.QProgressDialog()
+    self.pbar.canceled.connect(self.cancel)
+    self.pbar.rejected.connect(self.cancel)
+    self.pbar.hide()
+
+    self.timer = QtCore.QTimer()
+
+  def clean(self):
+    self.timer.stop()
+    try:
+      self.timer.timeout.disconnect()
+    except TypeError:
+      pass
+    try:
+      self.pbar.accepted.disconnect()
+    except TypeError:
+      pass
+
+  def cancel_delayed(self):
+    for delayed in self.delayed_queries:
+      log('match_action').info("async task cancelled: %s", repr(delayed))
+      delayed.cancel()
+    self.delayed_queries = []
+
+  def cancel(self):
+    log('match_action').info("match action cancelled")
+    self.clean()
+    self.cancel_delayed()
+
   @staticmethod
   def calc_file_version_hash():
     version_obj = {}
@@ -69,62 +100,55 @@ def response_handler(self, file_version):
     return True
 
   def start_upload(self):
+    log('match_action').info("Data upload started")
+
     self.functions = set(idautils.Functions())
 
-    self.pbar = QtWidgets.QProgressDialog()
     self.pbar.setLabelText("Processing IDB... You may continue working,\nbut "
                            "please avoid making any ground-breaking changes.")
     self.pbar.setRange(0, len(self.functions))
     self.pbar.setValue(0)
-    self.pbar.canceled.connect(self.cancel_upload)
-    self.pbar.rejected.connect(self.reject_upload)
     self.pbar.accepted.connect(self.accept_upload)
+    self.pbar.show()
 
-    self.timer = QtCore.QTimer()
     self.timer.timeout.connect(self.perform_upload)
     self.timer.start(0)
 
     return True
 
   def perform_upload(self):
-    try:
-      offset = self.functions.pop()
-    except KeyError:
-      self.timer.stop()
+    if not self.functions:
       return
 
-    try:
-      func = instances.FunctionInstance(self.file_version_id, offset)
-      self.instance_set.append(func.serialize())
-
-      if len(self.instance_set) >= 100:
-        network.delayed_query("POST", "collab/instances/",
-                              params=self.instance_set, json=True,
-                              callback=self.progress_advance)
-        self.instance_set = []
-        self.pbar.setMaximum(self.pbar.maximum() + 1)
-      self.progress_advance()
-    except Exception:
-      self.cancel_upload()
-      raise
+    # pop a function, serialize and add to the ready set
+    offset = self.functions.pop()
+    func = instances.FunctionInstance(self.file_version_id, offset)
+    self.instance_set.append(func.serialize())
+
+    # if ready set contains 100 or more functions, or if we just poped the last
+    # function clear and upload entire ready set to the server.
+    if len(self.instance_set) >= 100 or not self.functions:
+      q = network.QueryWorker("POST", "collab/instances/",
+                              params=self.instance_set, json=True)
+      q.start(self.progress_advance)
+      self.instance_set = []
+      self.pbar.setMaximum(self.pbar.maximum() + 1)
+    self.progress_advance()
 
   def progress_advance(self, result=None):
     del result
     new_value = self.pbar.value() + 1
-    self.pbar.setValue(new_value)
     if new_value >= self.pbar.maximum():
       self.pbar.accept()
+    else:
+      self.pbar.setValue(new_value)
 
-  def cancel_upload(self):
-    self.timer.stop()
-    self.timer = None
-    self.pbar = None
+  def accept_upload(self):
+    log('match_action').info("Data upload completed successfully")
 
-  def reject_upload(self):
-    self.cancel_upload()
+    self.clean()
+    self.delayed_queries = []
 
-  def accept_upload(self):
-    self.cancel_upload()
     self.start_task()
 
   def start_task(self):
@@ -148,19 +172,15 @@ def start_task(self):
     r = network.query("POST", "collab/tasks/", params=params, json=True)
     self.task_id = r['id']
 
-    self.pbar = QtWidgets.QProgressDialog()
     self.pbar.setLabelText("Waiting for remote matching... You may continue "
                            "working without any limitations.")
     self.pbar.setRange(0, int(r['progress_max']) if r['progress_max'] else 0)
     self.pbar.setValue(int(r['progress']))
-    self.pbar.canceled.connect(self.cancel_task)
-    self.pbar.rejected.connect(self.reject_task)
     self.pbar.accepted.connect(self.accept_task)
     self.pbar.show()
 
-    self.timer = QtCore.QTimer()
     self.timer.timeout.connect(self.perform_task)
-    self.timer.start(1000)
+    self.timer.start(200)
 
   def perform_task(self):
     try:
@@ -171,24 +191,91 @@ def perform_task(self):
       progress = int(r['progress'])
       status = r['status']
       if status == 'failed':
-        self.pbar.reject()
+        self.pbar.cancel()
       elif progress_max:
         self.pbar.setMaximum(progress_max)
         if progress >= progress_max:
           self.pbar.accept()
         else:
           self.pbar.setValue(progress)
     except Exception:
-      self.cancel_task()
+      self.cancel()
+      log('match_action').exception("perform update failed")
       raise
 
-  def cancel_task(self):
-    self.timer.stop()
-    self.timer = None
-    self.pbar = None
+  def accept_task(self):
+    log('match_action').info("Remote task completed successfully")
 
-  def reject_task(self):
-    self.cancel_task()
+    self.clean()
+    self.delayed_queries = []
 
-  def accept_task(self):
-    self.cancel_task()
+    self.start_results()
+
+  def start_results(self):
+    self.pbar.setLabelText("Receiving match results...")
+    self.pbar.setRange(0, 0)
+    self.pbar.setValue(0)
+    self.pbar.accepted.connect(self.accept_results)
+    self.pbar.show()
+
+    self.results = MatchResultDialog(self.task_id)
+
+    log('match_action').info("Result download started")
+    locals_url = "collab/tasks/{}/locals/".format(self.task_id)
+    q = network.QueryWorker("GET", locals_url, json=True, paginate=True,
+                            params={'limit': 100})
+    q.start(self.handle_locals)
+    self.delayed_queries.append(q)
+
+    remotes_url = "collab/tasks/{}/remotes/".format(self.task_id)
+    q = network.QueryWorker("GET", remotes_url, json=True, paginate=True,
+                            params={'limit': 100})
+    q.start(self.handle_remotes)
+    self.delayed_queries.append(q)
+
+    matches_url = "collab/tasks/{}/matches/".format(self.task_id)
+    q = network.QueryWorker("GET", matches_url, json=True, paginate=True,
+                            params={'limit': 100})
+    q.start(self.handle_matches)
+    self.delayed_queries.append(q)
+
+  def handle_locals(self, response):
+    new_locals = {obj['id']: obj for obj in response['results']}
+    self.results.add_locals(new_locals)
+
+    self.handle_page(response)
+
+  def handle_remotes(self, response):
+    new_remotes = {obj['id']: obj for obj in response['results']}
+    self.results.add_remotes(new_remotes)
+
+    self.handle_page(response)
+
+  def handle_matches(self, response):
+    def rename(o):
+      o['local_id'] = o.pop('from_instance')
+      o['remote_id'] = o.pop('to_instance')
+      return o
+
+    new_matches = map(rename, response['results'])
+    self.results.add_matches(new_matches)
+
+    self.handle_page(response)
+
+  def handle_page(self, response):
+    if 'previous' not in response or not response['previous']:
+      self.pbar.setMaximum(self.pbar.maximum() + response['count'])
+
+    new_value = max(self.pbar.value(), 0) + len(response['results'])
+    if new_value >= self.pbar.maximum():
+      self.pbar.accept()
+    else:
+      self.pbar.setValue(new_value)
+
+  def accept_results(self):
+    log('match_action').info("Result download completed successfully")
+
+    self.clean()
+    self.delayed_queries = []
+
+    self.results.show()
@@ -3,5 +3,6 @@
 import match
 import project
 import settings
+import serializedgraph
 
-__all__ = [base, login, match, project, settings]
+__all__ = ['base', 'login', 'match', 'project', 'settings', 'serializedgraph']
@@ -7,10 +7,11 @@
 
 
 class BaseDialog(QtWidgets.QDialog):
-  def __init__(self, title="", reject_handler=None, submit_handler=None,
-               response_handler=None, exception_handler=None, **kwargs):
+  def __init__(self, title="", modal=True, reject_handler=None,
+               submit_handler=None, response_handler=None,
+               exception_handler=None, **kwargs):
     super(BaseDialog, self).__init__(**kwargs)
-    self.setModal(True)
+    self.setModal(modal)
     self.setWindowTitle(title)
     self.reject_handler = reject_handler
     self.submit_handler = submit_handler
@@ -61,8 +62,7 @@ def submit_base(self):
       return
 
     # if received a query_worker, execute it and handle response
-    network.delayed_worker(query_worker, self.response_base,
-                           self.exception_base)
+    query_worker.start(self.response_base, self.exception_base)
 
   def reject_base(self):
     if self.reject_handler: