Move Pyrit logic to tools/pyrit.py

wifite/model/handshake.py

@@ -4,6 +4,7 @@
 from ..util.process import Process
 from ..util.color import Color
 from ..tools.tshark import Tshark
+from ..tools.pyrit import Pyrit
 
 import re, os
 
@@ -67,9 +68,6 @@ class Handshake(object):
         return False
 
 
-    def tshark_bssid_essid_pairs(self):
-        '''Returns list of tuples: (bssid,essid) found in capfile'''
-
     def tshark_handshakes(self):
         ''' Returns True if tshark identifies a handshake, False otherwise '''
         tshark_bssids = Tshark.bssids_with_handshakes(self.capfile, bssid=self.bssid)
@@ -98,54 +96,9 @@ class Handshake(object):
         return []
 
 
-    def pyrit_command(self):
-        return [
-            'pyrit',
-            '-r', self.capfile,
-            'analyze'
-        ]
-
     def pyrit_handshakes(self):
-        ''' Returns True if pyrit identifies a handshake, False otherwise '''
-        if not Process.exists('pyrit'):
-            return []
-
-        bssid_essid_pairs = set()
-        hit_target = False
-        current_bssid = self.bssid
-        current_essid = self.essid
-        proc = Process(self.pyrit_command(), devnull=False)
-        for line in proc.stdout().split('\n'):
-            mac_regex = ('[a-zA-Z0-9]{2}:' * 6)[:-1]
-            match = re.search("^#\d+: AccessPoint (%s) \('(.*)'\):$"
-                % (mac_regex), line)
-            if match:
-                # We found a BSSID and ESSID
-                (bssid, essid) = match.groups()
-
-                # Compare to what we're searching for
-                if self.bssid and self.bssid.lower() == bssid.lower():
-                    current_essid = essid
-                    hit_target = True
-                    continue
-
-                elif self.essid and self.essid == essid:
-                    current_bssid = bssid
-                    hit_target = True
-                    continue
-
-                elif not self.bssid and not self.essid:
-                    # We don't know either
-                    current_bssid = bssid
-                    current_essid = essid
-                    hit_target = True
-                else:
-                    hit_Target = False # This AccessPoint is not what we're looking for
-            else:
-                # Line does not contain AccessPoint
-                if hit_target and ', good' in line:
-                    bssid_essid_pairs.add( (current_bssid, current_essid) )
-        return [x for x in bssid_essid_pairs]
+        ''' Returns list of BSSID,ESSID tuples if pyrit identifies a handshake'''
+        return Pyrit.bssid_essid_with_handshakes(self.capfile, bssid=self.bssid, essid=self.essid)
 
 
     def aircrack_handshakes(self):

wifite/tools/pyrit.py

@@ -0,0 +1,59 @@
+#!/usr/bin/python2.7
+# -*- coding: utf-8 -*-
+
+from ..util.process import Process
+import re
+
+class Pyrit(object):
+    ''' Wrapper for Pyrit program. '''
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def exists():
+        return Process.exists('pyrit')
+
+    @staticmethod
+    def bssid_essid_with_handshakes(capfile, bssid=None, essid=None):
+        if not Pyrit.exists():
+            return []
+
+        command = [
+            'pyrit',
+            '-r', capfile,
+            'analyze'
+        ]
+        pyrit = Process(command, devnull=False)
+
+        current_bssid = current_essid = None
+        bssid_essid_pairs = set()
+
+        '''
+        #1: AccessPoint 18:a6:f7:31:d2:06 ('TP-LINK_D206'):
+          #1: Station 08:66:98:b2:ab:28, 1 handshake(s):
+              #1: HMAC_SHA1_AES, good, spread 1
+                #2: Station ac:63:be:3a:a2:f4
+        '''
+
+        for line in pyrit.stdout().split('\n'):
+            mac_regex = ('[a-zA-Z0-9]{2}:' * 6)[:-1]
+            match = re.search("^#\d+: AccessPoint (%s) \('(.*)'\):$" % (mac_regex), line)
+            if match:
+                # We found a new BSSID and ESSID
+                (current_bssid, current_essid) = match.groups()
+
+                if bssid is not None and bssid.lower() != current_bssid:
+                    current_bssid = None
+                    current_essid = None
+                elif essid is not None and essid != current_essid:
+                    current_bssid = None
+                    current_essid = None
+
+            elif current_bssid is not None and current_essid is not None:
+                # We hit an AP that we care about.
+                # Line does not contain AccessPoint, see if it's "good"
+                if ', good' in line:
+                    bssid_essid_pairs.add( (current_bssid, current_essid) )
+
+        return list(bssid_essid_pairs)