fix: supporting channels greater than 140 for 5g (closes #583)

evilsocket · Nov 13, 2019 · a8c07ba · a8c07ba
1 parent 9dcc647
commit a8c07ba
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 28 deletions.
diff --git a/pwnagotchi/ai/featurizer.py b/pwnagotchi/ai/featurizer.py
@@ -4,31 +4,37 @@
 
 MAX_EPOCH_DURATION = 1024
 
-histogram_size = wifi.NumChannels
 
-shape = (1,
-         # aps per channel
-         histogram_size +
-         # clients per channel
-         histogram_size +
-         # peers per channel
-         histogram_size +
-         # duration
-         1 +
-         # inactive
-         1 +
-         # active
-         1 +
-         # missed
-         1 +
-         # hops
-         1 +
-         # deauths
-         1 +
-         # assocs
-         1 +
-         # handshakes
-         1)
+def describe(extended=False):
+    if not extended:
+        histogram_size = wifi.NumChannels
+    else:
+        # see https://github.com/evilsocket/pwnagotchi/issues/583
+        histogram_size = wifi.NumChannelsExt
+
+    return histogram_size, (1,
+                            # aps per channel
+                            histogram_size +
+                            # clients per channel
+                            histogram_size +
+                            # peers per channel
+                            histogram_size +
+                            # duration
+                            1 +
+                            # inactive
+                            1 +
+                            # active
+                            1 +
+                            # missed
+                            1 +
+                            # hops
+                            1 +
+                            # deauths
+                            1 +
+                            # assocs
+                            1 +
+                            # handshakes
+                            1)
 
 
 def featurize(state, step):

diff --git a/pwnagotchi/ai/gym.py b/pwnagotchi/ai/gym.py
@@ -34,10 +34,14 @@ def __init__(self, agent, epoch):
         self._epoch_num = 0
         self._last_render = None
 
-        channels = agent.supported_channels()
+        # see https://github.com/evilsocket/pwnagotchi/issues/583
+        self._supported_channels = agent.supported_channels()
+        self._extended_spectrum = any(ch > 140 for ch in self._supported_channels)
+        self._histogram_size, self._observation_shape = featurizer.describe(self._extended_spectrum)
+
         Environment.params += [
             Parameter('_channel_%d' % ch, min_value=0, max_value=1, meta=ch + 1) for ch in
-            range(featurizer.histogram_size) if ch + 1 in channels
+            range(self._histogram_size) if ch + 1 in self._supported_channels
         ]
 
         self.last = {
@@ -50,7 +54,7 @@ def __init__(self, agent, epoch):
         }
 
         self.action_space = spaces.MultiDiscrete([p.space_size() for p in Environment.params if p.trainable])
-        self.observation_space = spaces.Box(low=0, high=1, shape=featurizer.shape, dtype=np.float32)
+        self.observation_space = spaces.Box(low=0, high=1, shape=self._observation_shape, dtype=np.float32)
         self.reward_range = reward.range
 
     @staticmethod
@@ -118,7 +122,7 @@ def reset(self):
         return self.last['state_v']
 
     def _render_histogram(self, hist):
-        for ch in range(featurizer.histogram_size):
+        for ch in range(self._histogram_size):
             if hist[ch]:
                 logging.info("      CH %d: %s" % (ch + 1, hist[ch]))
 

diff --git a/pwnagotchi/mesh/wifi.py b/pwnagotchi/mesh/wifi.py
@@ -1,4 +1,5 @@
 NumChannels = 140
+NumChannelsExt = 165  # see https://github.com/evilsocket/pwnagotchi/issues/583
 
 
 def freq_to_channel(freq):