From 55d99836e7e74a36c6d2ea4306bcb8ab3046f432 Mon Sep 17 00:00:00 2001
From: Simone Margaritelli <evilsocket@gmail.com>
Date: Tue, 8 Oct 2019 14:54:03 +0200
Subject: [PATCH] fix: on_internet_available plugins callback is now called for
 both MANU and AUTO mode (fixes #210)

---
 bin/pwnagotchi                                | 24 ++++---
 pwnagotchi/agent.py                           |  5 ++
 pwnagotchi/log.py                             | 68 ++++++++++---------
 pwnagotchi/mesh/utils.py                      |  3 +
 pwnagotchi/plugins/default/auto-backup.py     |  4 +-
 pwnagotchi/plugins/default/auto-update.py     |  4 +-
 pwnagotchi/plugins/default/example.py         |  2 +-
 pwnagotchi/plugins/default/grid.py            | 54 ++++++++-------
 pwnagotchi/plugins/default/onlinehashcrack.py |  5 +-
 pwnagotchi/plugins/default/twitter.py         | 22 +++---
 pwnagotchi/plugins/default/wigle.py           |  5 +-
 pwnagotchi/plugins/default/wpa-sec.py         |  5 +-
 pwnagotchi/ui/view.py                         | 16 ++---
 pwnagotchi/voice.py                           | 24 +++----
 14 files changed, 139 insertions(+), 102 deletions(-)

diff --git a/bin/pwnagotchi b/bin/pwnagotchi
index f72d84c..9cd77d8 100755
--- a/bin/pwnagotchi
+++ b/bin/pwnagotchi
@@ -9,7 +9,6 @@ if __name__ == '__main__':
     import pwnagotchi.utils as utils
     import pwnagotchi.plugins as plugins
 
-    from pwnagotchi.log import SessionParser
     from pwnagotchi.identity import KeyPair
     from pwnagotchi.agent import Agent
     from pwnagotchi.ui.display import Display
@@ -51,22 +50,23 @@ if __name__ == '__main__':
     elif args.do_manual:
         logging.info("entering manual mode ...")
 
-        log = SessionParser(config)
+        agent.last_session.parse()
+
         logging.info(
             "the last session lasted %s (%d completed epochs, trained for %d), average reward:%s (min:%s max:%s)" % (
-                log.duration_human,
-                log.epochs,
-                log.train_epochs,
-                log.avg_reward,
-                log.min_reward,
-                log.max_reward))
+                agent.last_session.duration_human,
+                agent.last_session.epochs,
+                agent.last_session.train_epochs,
+                agent.last_session.avg_reward,
+                agent.last_session.min_reward,
+                agent.last_session.max_reward))
 
         while True:
-            display.on_manual_mode(log)
+            display.on_manual_mode(agent.last_session)
             time.sleep(1)
 
             if Agent.is_connected():
-                plugins.on('internet_available', display, keypair, config, log)
+                plugins.on('internet_available', agent)
 
     else:
         logging.info("entering auto mode ...")
@@ -104,5 +104,9 @@ if __name__ == '__main__':
                 # WiFi electromagnetic fields affect time like gravitational fields
                 # affect ours ... neat ^_^
                 agent.next_epoch()
+
+                if Agent.is_connected():
+                    plugins.on('internet_available', agent)
+
             except Exception as e:
                 logging.exception("main loop exception")
diff --git a/pwnagotchi/agent.py b/pwnagotchi/agent.py
index 71e87f5..b29e164 100644
--- a/pwnagotchi/agent.py
+++ b/pwnagotchi/agent.py
@@ -9,6 +9,7 @@ import _thread
 
 import pwnagotchi.utils as utils
 import pwnagotchi.plugins as plugins
+from pwnagotchi.log import LastSession
 from pwnagotchi.bettercap import Client
 from pwnagotchi.mesh.utils import AsyncAdvertiser
 from pwnagotchi.ai.train import AsyncTrainer
@@ -35,6 +36,7 @@ class Agent(Client, AsyncAdvertiser, AsyncTrainer):
         self._last_pwnd = None
         self._history = {}
         self._handshakes = {}
+        self.last_session = LastSession(self._config)
 
     @staticmethod
     def is_connected():
@@ -48,6 +50,9 @@ class Agent(Client, AsyncAdvertiser, AsyncTrainer):
     def config(self):
         return self._config
 
+    def view(self):
+        return self._view
+
     def supported_channels(self):
         return self._supported_channels
 
diff --git a/pwnagotchi/log.py b/pwnagotchi/log.py
index dab58e4..0f05789 100644
--- a/pwnagotchi/log.py
+++ b/pwnagotchi/log.py
@@ -11,7 +11,7 @@ from file_read_backwards import FileReadBackwards
 LAST_SESSION_FILE = '/root/.pwnagotchi-last-session'
 
 
-class SessionParser(object):
+class LastSession(object):
     EPOCH_TOKEN = '[epoch '
     EPOCH_PARSER = re.compile(r'^.+\[epoch (\d+)\] (.+)')
     EPOCH_DATA_PARSER = re.compile(r'([a-z_]+)=([^\s]+)')
@@ -70,27 +70,27 @@ class SessionParser(object):
             if started_at is None:
                 started_at = stopped_at
 
-            if SessionParser.DEAUTH_TOKEN in line and line not in cache:
+            if LastSession.DEAUTH_TOKEN in line and line not in cache:
                 self.deauthed += 1
                 cache[line] = 1
 
-            elif SessionParser.ASSOC_TOKEN in line and line not in cache:
+            elif LastSession.ASSOC_TOKEN in line and line not in cache:
                 self.associated += 1
                 cache[line] = 1
 
-            elif SessionParser.HANDSHAKE_TOKEN in line and line not in cache:
+            elif LastSession.HANDSHAKE_TOKEN in line and line not in cache:
                 self.handshakes += 1
                 cache[line] = 1
 
-            elif SessionParser.TRAINING_TOKEN in line:
+            elif LastSession.TRAINING_TOKEN in line:
                 self.train_epochs += 1
 
-            elif SessionParser.EPOCH_TOKEN in line:
+            elif LastSession.EPOCH_TOKEN in line:
                 self.epochs += 1
-                m = SessionParser.EPOCH_PARSER.findall(line)
+                m = LastSession.EPOCH_PARSER.findall(line)
                 if m:
                     epoch_num, epoch_data = m[0]
-                    m = SessionParser.EPOCH_DATA_PARSER.findall(epoch_data)
+                    m = LastSession.EPOCH_DATA_PARSER.findall(epoch_data)
                     for key, value in m:
                         if key == 'reward':
                             reward = float(value)
@@ -101,7 +101,7 @@ class SessionParser(object):
                             elif reward > self.max_reward:
                                 self.max_reward = reward
 
-            elif SessionParser.PEER_TOKEN in line:
+            elif LastSession.PEER_TOKEN in line:
                 m = self._peer_parser.findall(line)
                 if m:
                     name, pubkey, rssi, sid, pwnd_tot, uptime = m[0]
@@ -134,6 +134,30 @@ class SessionParser(object):
         self.duration_human = ', '.join(self.duration_human)
         self.avg_reward /= (self.epochs if self.epochs else 1)
 
+    def parse(self):
+        lines = []
+
+        if os.path.exists(self.path):
+            with FileReadBackwards(self.path, encoding="utf-8") as fp:
+                for line in fp:
+                    line = line.strip()
+                    if line != "" and line[0] != '[':
+                        continue
+                    lines.append(line)
+                    if LastSession.START_TOKEN in line:
+                        break
+            lines.reverse()
+
+        if len(lines) == 0:
+            lines.append("Initial Session");
+
+        self.last_session = lines
+        self.last_session_id = hashlib.md5(lines[0].encode()).hexdigest()
+        self.last_saved_session_id = self._get_last_saved_session_id()
+
+        self._parse_stats()
+        self.parsed = True
+
     def __init__(self, config):
         self.config = config
         self.voice = Voice(lang=config['main']['lang'])
@@ -150,28 +174,10 @@ class SessionParser(object):
         self.last_peer = None
         self._peer_parser = re.compile(
             'detected unit (.+)@(.+) \(v.+\) on channel \d+ \(([\d\-]+) dBm\) \[sid:(.+) pwnd_tot:(\d+) uptime:(\d+)\]')
-
-        lines = []
-
-        if os.path.exists(self.path):
-            with FileReadBackwards(self.path, encoding="utf-8") as fp:
-                for line in fp:
-                    line = line.strip()
-                    if line != "" and line[0] != '[':
-                        continue
-                    lines.append(line)
-                    if SessionParser.START_TOKEN in line:
-                        break
-            lines.reverse()
-
-        if len(lines) == 0:
-            lines.append("Initial Session");
-
-        self.last_session = lines
-        self.last_session_id = hashlib.md5(lines[0].encode()).hexdigest()
-        self.last_saved_session_id = self._get_last_saved_session_id()
-
-        self._parse_stats()
+        self.last_session = []
+        self.last_session_id = None
+        self.last_saved_session_id = None
+        self.parsed = False
 
     def is_new(self):
         return self.last_session_id != self.last_saved_session_id
diff --git a/pwnagotchi/mesh/utils.py b/pwnagotchi/mesh/utils.py
index a775ad4..3a67ed9 100644
--- a/pwnagotchi/mesh/utils.py
+++ b/pwnagotchi/mesh/utils.py
@@ -12,6 +12,9 @@ class AsyncAdvertiser(object):
         self._keypair = keypair
         self._advertiser = None
 
+    def keypair(self):
+        return self._keypair
+
     def start_advertising(self):
         _thread.start_new_thread(self._adv_worker, ())
 
diff --git a/pwnagotchi/plugins/default/auto-backup.py b/pwnagotchi/plugins/default/auto-backup.py
index 48d8597..7c9901c 100644
--- a/pwnagotchi/plugins/default/auto-backup.py
+++ b/pwnagotchi/plugins/default/auto-backup.py
@@ -33,7 +33,7 @@ def on_loaded():
     logging.info("AUTO-BACKUP: Successfuly loaded.")
 
 
-def on_internet_available(display, keypair, config, log):
+def on_internet_available(agent):
     global STATUS
 
     if READY:
@@ -42,6 +42,8 @@ def on_internet_available(display, keypair, config, log):
 
         files_to_backup = " ".join(OPTIONS['files'])
         try:
+            display = agent.view()
+
             logging.info("AUTO-BACKUP: Backing up ...")
             display.set('status', 'Backing up ...')
             display.update()
diff --git a/pwnagotchi/plugins/default/auto-update.py b/pwnagotchi/plugins/default/auto-update.py
index dab44b3..8a01a45 100644
--- a/pwnagotchi/plugins/default/auto-update.py
+++ b/pwnagotchi/plugins/default/auto-update.py
@@ -23,13 +23,15 @@ def on_loaded():
     READY = True
 
 
-def on_internet_available(display, keypair, config, log):
+def on_internet_available(agent):
     global STATUS
 
     if READY:
         if STATUS.newer_then_days(OPTIONS['interval']):
             return
 
+        display = agent.view()
+
         try:
             display.set('status', 'Updating ...')
             display.update()
diff --git a/pwnagotchi/plugins/default/example.py b/pwnagotchi/plugins/default/example.py
index 2a82a18..3f3e570 100644
--- a/pwnagotchi/plugins/default/example.py
+++ b/pwnagotchi/plugins/default/example.py
@@ -20,7 +20,7 @@ def on_loaded():
 
 
 # called in manual mode when there's internet connectivity
-def on_internet_available(ui, keypair, config, log):
+def on_internet_available(agent):
     pass
 
 
diff --git a/pwnagotchi/plugins/default/grid.py b/pwnagotchi/plugins/default/grid.py
index dfc54ad..5427dbb 100644
--- a/pwnagotchi/plugins/default/grid.py
+++ b/pwnagotchi/plugins/default/grid.py
@@ -22,16 +22,16 @@ def on_loaded():
     logging.info("api plugin loaded.")
 
 
-def get_api_token(log, keys):
+def get_api_token(last_session, keys):
     global AUTH
 
     if AUTH.newer_then_minutes(25) and AUTH.data is not None and 'token' in AUTH.data:
         return AUTH.data['token']
 
     if AUTH.data is None:
-        logging.info("api: enrolling unit ...")
+        logging.info("grid: enrolling unit ...")
     else:
-        logging.info("api: refreshing token ...")
+        logging.info("grid: refreshing token ...")
 
     identity = "%s@%s" % (pwnagotchi.name(), keys.fingerprint)
     # sign the identity string to prove we own both keys
@@ -43,16 +43,16 @@ def get_api_token(log, keys):
         'public_key': keys.pub_key_pem_b64,
         'signature': signature_b64,
         'data': {
-            'duration': log.duration,
-            'epochs': log.epochs,
-            'train_epochs': log.train_epochs,
-            'avg_reward': log.avg_reward,
-            'min_reward': log.min_reward,
-            'max_reward': log.max_reward,
-            'deauthed': log.deauthed,
-            'associated': log.associated,
-            'handshakes': log.handshakes,
-            'peers': log.peers,
+            'duration': last_session.duration,
+            'epochs': last_session.epochs,
+            'train_epochs': last_session.train_epochs,
+            'avg_reward': last_session.avg_reward,
+            'min_reward': last_session.min_reward,
+            'max_reward': last_session.max_reward,
+            'deauthed': last_session.deauthed,
+            'associated': last_session.associated,
+            'handshakes': last_session.handshakes,
+            'peers': last_session.peers,
             'uname': subprocess.getoutput("uname -a")
         }
     }
@@ -63,13 +63,13 @@ def get_api_token(log, keys):
 
     AUTH.update(data=r.json())
 
-    logging.info("api: done")
+    logging.info("grid: done")
 
     return AUTH.data["token"]
 
 
 def parse_pcap(filename):
-    logging.info("api: parsing %s ..." % filename)
+    logging.info("grid: parsing %s ..." % filename)
 
     net_id = os.path.basename(filename).replace('.pcap', '')
 
@@ -91,15 +91,15 @@ def parse_pcap(filename):
     try:
         info = extract_from_pcap(filename, [WifiInfo.BSSID, WifiInfo.ESSID])
     except Exception as e:
-        logging.error("api: %s" % e)
+        logging.error("grid: %s" % e)
 
     return info[WifiInfo.ESSID], info[WifiInfo.BSSID]
 
 
-def api_report_ap(log, keys, token, essid, bssid):
+def api_report_ap(last_session, keys, token, essid, bssid):
     while True:
         token = AUTH.data['token']
-        logging.info("api: reporting %s (%s)" % (essid, bssid))
+        logging.info("grid: reporting %s (%s)" % (essid, bssid))
         try:
             api_address = 'https://api.pwnagotchi.ai/api/v1/unit/report/ap'
             headers = {'Authorization': 'access_token %s' % token}
@@ -111,21 +111,23 @@ def api_report_ap(log, keys, token, essid, bssid):
             if r.status_code != 200:
                 if r.status_code == 401:
                     logging.warning("token expired")
-                    token = get_api_token(log, keys)
+                    token = get_api_token(last_session, keys)
                     continue
                 else:
                     raise Exception("(status %d) %s" % (r.status_code, r.text))
             else:
                 return True
         except Exception as e:
-            logging.error("api: %s" % e)
+            logging.error("grid: %s" % e)
             return False
 
 
-def on_internet_available(ui, keys, config, log):
+def on_internet_available(agent):
     global REPORT
 
     try:
+        config = agent.config()
+        keys = agent.keypair()
 
         pcap_files = glob.glob(os.path.join(config['bettercap']['handshakes'], "*.pcap"))
         num_networks = len(pcap_files)
@@ -134,10 +136,10 @@ def on_internet_available(ui, keys, config, log):
         num_new = num_networks - num_reported
 
         if num_new > 0:
-            logging.info("api: %d new networks to report" % num_new)
-            token = get_api_token(log, keys)
-
             if OPTIONS['report']:
+                logging.info("grid: %d new networks to report" % num_new)
+                token = get_api_token(agent.last_session, agent.keypair())
+
                 for pcap_file in pcap_files:
                     net_id = os.path.basename(pcap_file).replace('.pcap', '')
                     do_skip = False
@@ -151,11 +153,11 @@ def on_internet_available(ui, keys, config, log):
                     if net_id not in reported and not do_skip:
                         essid, bssid = parse_pcap(pcap_file)
                         if bssid:
-                            if api_report_ap(log, keys, token, essid, bssid):
+                            if api_report_ap(agent.last_session, keys, token, essid, bssid):
                                 reported.append(net_id)
                                 REPORT.update(data={'reported': reported})
             else:
-                logging.info("api: reporting disabled")
+                logging.debug("grid: reporting disabled")
 
     except Exception as e:
         logging.exception("error while enrolling the unit")
diff --git a/pwnagotchi/plugins/default/onlinehashcrack.py b/pwnagotchi/plugins/default/onlinehashcrack.py
index 9771b88..25a18d6 100644
--- a/pwnagotchi/plugins/default/onlinehashcrack.py
+++ b/pwnagotchi/plugins/default/onlinehashcrack.py
@@ -55,11 +55,14 @@ def _upload_to_ohc(path, timeout=30):
             raise e
 
 
-def on_internet_available(display, keypair, config, log):
+def on_internet_available(agent):
     """
     Called in manual mode when there's internet connectivity
     """
     if READY:
+        display = agent.view()
+        config = agent.config()
+
         handshake_dir = config['bettercap']['handshakes']
         handshake_filenames = os.listdir(handshake_dir)
         handshake_paths = [os.path.join(handshake_dir, filename) for filename in handshake_filenames if filename.endswith('.pcap')]
diff --git a/pwnagotchi/plugins/default/twitter.py b/pwnagotchi/plugins/default/twitter.py
index 8f21f25..fd247ca 100644
--- a/pwnagotchi/plugins/default/twitter.py
+++ b/pwnagotchi/plugins/default/twitter.py
@@ -14,8 +14,12 @@ def on_loaded():
 
 
 # called in manual mode when there's internet connectivity
-def on_internet_available(ui, keypair, config, log):
-    if log.is_new() and log.handshakes > 0:
+def on_internet_available(agent):
+    config = agent.config()
+    display = agent.view()
+    last_session = agent.last_session
+
+    if last_session.is_new() and last_session.handshakes > 0:
         try:
             import tweepy
         except ImportError:
@@ -26,20 +30,20 @@ def on_internet_available(ui, keypair, config, log):
 
         picture = '/dev/shm/pwnagotchi.png'
 
-        ui.on_manual_mode(log)
-        ui.update(force=True)
-        ui.image().save(picture, 'png')
-        ui.set('status', 'Tweeting...')
-        ui.update(force=True)
+        display.on_manual_mode(last_session)
+        display.update(force=True)
+        display.image().save(picture, 'png')
+        display.set('status', 'Tweeting...')
+        display.update(force=True)
 
         try:
             auth = tweepy.OAuthHandler(OPTIONS['consumer_key'], OPTIONS['consumer_secret'])
             auth.set_access_token(OPTIONS['access_token_key'], OPTIONS['access_token_secret'])
             api = tweepy.API(auth)
 
-            tweet = Voice(lang=config['main']['lang']).on_log_tweet(log)
+            tweet = Voice(lang=config['main']['lang']).on_last_session_tweet(last_session)
             api.update_with_media(filename=picture, status=tweet)
-            log.save_session_id()
+            last_session.save_session_id()
 
             logging.info("tweeted: %s" % tweet)
         except Exception as e:
diff --git a/pwnagotchi/plugins/default/wigle.py b/pwnagotchi/plugins/default/wigle.py
index 4f89ad4..c72c1de 100644
--- a/pwnagotchi/plugins/default/wigle.py
+++ b/pwnagotchi/plugins/default/wigle.py
@@ -196,7 +196,7 @@ def _send_to_wigle(lines, api_key, timeout=30):
         raise re_e
 
 
-def on_internet_available(display, keypair, config, log):
+def on_internet_available(agent):
     from scapy.all import RadioTap, Dot11Elt, Dot11Beacon, rdpcap, Scapy_Exception, Dot11, Dot11ProbeResp, Dot11AssoReq, \
         Dot11ReassoReq, Dot11EltRSN, Dot11EltVendorSpecific, Dot11EltMicrosoftWPA
     """
@@ -206,6 +206,9 @@ def on_internet_available(display, keypair, config, log):
     global SKIP
 
     if READY:
+        config = agent.config()
+        display = agent.view()
+
         handshake_dir = config['bettercap']['handshakes']
         all_files = os.listdir(handshake_dir)
         all_gps_files = [os.path.join(handshake_dir, filename)
diff --git a/pwnagotchi/plugins/default/wpa-sec.py b/pwnagotchi/plugins/default/wpa-sec.py
index 7d61b7b..0f6b854 100644
--- a/pwnagotchi/plugins/default/wpa-sec.py
+++ b/pwnagotchi/plugins/default/wpa-sec.py
@@ -54,11 +54,14 @@ def _upload_to_wpasec(path, timeout=30):
             raise e
 
 
-def on_internet_available(display, keypair, config, log):
+def on_internet_available(agent):
     """
     Called in manual mode when there's internet connectivity
     """
     if READY:
+        config = agent.config()
+        display = agent.view()
+
         handshake_dir = config['bettercap']['handshakes']
         handshake_filenames = os.listdir(handshake_dir)
         handshake_paths = [os.path.join(handshake_dir, filename) for filename in handshake_filenames if filename.endswith('.pcap')]
diff --git a/pwnagotchi/ui/view.py b/pwnagotchi/ui/view.py
index 3d247f9..354d910 100644
--- a/pwnagotchi/ui/view.py
+++ b/pwnagotchi/ui/view.py
@@ -163,17 +163,17 @@ class View(object):
         self.set('status', self._voice.on_ai_ready())
         self.update()
 
-    def on_manual_mode(self, log):
+    def on_manual_mode(self, last_session):
         self.set('mode', 'MANU')
-        self.set('face', faces.SAD if log.handshakes == 0 else faces.HAPPY)
-        self.set('status', self._voice.on_log(log))
-        self.set('epoch', "%04d" % log.epochs)
-        self.set('uptime', log.duration)
+        self.set('face', faces.SAD if last_session.handshakes == 0 else faces.HAPPY)
+        self.set('status', self._voice.on_last_session_data(last_session))
+        self.set('epoch', "%04d" % last_session.epochs)
+        self.set('uptime', last_session.duration)
         self.set('channel', '-')
-        self.set('aps', "%d" % log.associated)
-        self.set('shakes', '%d (%s)' % (log.handshakes, \
+        self.set('aps', "%d" % last_session.associated)
+        self.set('shakes', '%d (%s)' % (last_session.handshakes, \
                                         utils.total_unique_handshakes(self._config['bettercap']['handshakes'])))
-        self.set_closest_peer(log.last_peer, log.peers)
+        self.set_closest_peer(last_session.last_peer, last_session.peers)
 
     def is_normal(self):
         return self._state.get('face') not in (
diff --git a/pwnagotchi/voice.py b/pwnagotchi/voice.py
index aa12334..a5305f6 100644
--- a/pwnagotchi/voice.py
+++ b/pwnagotchi/voice.py
@@ -125,23 +125,23 @@ class Voice:
     def on_rebooting(self):
         return self._("Ops, something went wrong ... Rebooting ...")
 
-    def on_log(self, log):
-        status = self._('Kicked {num} stations\n').format(num=log.deauthed)
-        status += self._('Made {num} new friends\n').format(num=log.associated)
-        status += self._('Got {num} handshakes\n').format(num=log.handshakes)
-        if log.peers == 1:
+    def on_last_session_data(self, last_session):
+        status = self._('Kicked {num} stations\n').format(num=last_session.deauthed)
+        status += self._('Made {num} new friends\n').format(num=last_session.associated)
+        status += self._('Got {num} handshakes\n').format(num=last_session.handshakes)
+        if last_session.peers == 1:
             status += self._('Met 1 peer')
-        elif log.peers > 0:
-            status += self._('Met {num} peers').format(num=log.peers)
+        elif last_session.peers > 0:
+            status += self._('Met {num} peers').format(num=last_session.peers)
         return status
 
-    def on_log_tweet(self, log):
+    def on_last_session_tweet(self, last_session):
         return self._(
             'I\'ve been pwning for {duration} and kicked {deauthed} clients! I\'ve also met {associated} new friends and ate {handshakes} handshakes! #pwnagotchi #pwnlog #pwnlife #hacktheplanet #skynet').format(
-            duration=log.duration_human,
-            deauthed=log.deauthed,
-            associated=log.associated,
-            handshakes=log.handshakes)
+            duration=last_session.duration_human,
+            deauthed=last_session.deauthed,
+            associated=last_session.associated,
+            handshakes=last_session.handshakes)
 
     def hhmmss(self, count, fmt):
         if count > 1: