diff --git a/pwnagotchi/ai/featurizer.py b/pwnagotchi/ai/featurizer.py index d8f27a74..b4630e2d 100644 --- a/pwnagotchi/ai/featurizer.py +++ b/pwnagotchi/ai/featurizer.py @@ -5,36 +5,31 @@ import pwnagotchi.mesh.wifi as wifi MAX_EPOCH_DURATION = 1024 -def describe(extended=False): - if not extended: - histogram_size = wifi.NumChannels - else: - # see https://github.com/evilsocket/pwnagotchi/issues/583 - histogram_size = wifi.NumChannelsExt +histogram_size = wifi.NumChannels - 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) +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 featurize(state, step): diff --git a/pwnagotchi/ai/gym.py b/pwnagotchi/ai/gym.py index c188d31d..a9800658 100644 --- a/pwnagotchi/ai/gym.py +++ b/pwnagotchi/ai/gym.py @@ -35,14 +35,11 @@ class Environment(gym.Env): self._epoch_num = 0 self._last_render = None - # see https://github.com/evilsocket/pwnagotchi/issues/583 - self._supported_channels = agent.supported_channels() - self._extended_spectrum = any(ch > 150 for ch in self._supported_channels) - self._histogram_size, self._observation_shape = featurizer.describe(self._extended_spectrum) + channels = agent.supported_channels() Environment.params += [ Parameter('_channel_%d' % ch, min_value=0, max_value=1, meta=ch + 1) for ch in - range(self._histogram_size) if ch + 1 in self._supported_channels + range(featurizer.histogram_size) if ch + 1 in channels ] self.last = { @@ -55,7 +52,7 @@ class Environment(gym.Env): } 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=self._observation_shape, dtype=np.float32) + self.observation_space = spaces.Box(low=0, high=1, shape=featurizer.shape, dtype=np.float32) self.reward_range = reward.range @staticmethod @@ -123,7 +120,7 @@ class Environment(gym.Env): return self.last['state_v'] def _render_histogram(self, hist): - for ch in range(self._histogram_size): + for ch in range(featurizer.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 index ea81022c..aa2a1da1 100644 --- a/pwnagotchi/mesh/wifi.py +++ b/pwnagotchi/mesh/wifi.py @@ -1,5 +1,4 @@ -NumChannels = 140 -NumChannelsExt = 165 # see https://github.com/evilsocket/pwnagotchi/issues/583 +NumChannels = 177 def freq_to_channel(freq): diff --git a/pwnagotchi/utils.py b/pwnagotchi/utils.py index 1a346f83..2cc25bc0 100644 --- a/pwnagotchi/utils.py +++ b/pwnagotchi/utils.py @@ -130,10 +130,10 @@ def merge_config(user, default): def keys_to_str(data): - if isinstance(data,list): + if isinstance(data, list): converted_list = list() for item in data: - if isinstance(item,list) or isinstance(item,dict): + if isinstance(item, list) or isinstance(item, dict): converted_list.append(keys_to_str(item)) else: converted_list.append(item) @@ -277,7 +277,8 @@ def load_config(args): elif config['ui']['display']['type'] in ('ws_154inch', 'ws154inch', 'waveshare_154inch', 'waveshare154inch'): config['ui']['display']['type'] = 'waveshare154inch' - elif config['ui']['display']['type'] in ('waveshare144lcd', 'ws_144inch', 'ws144inch', 'waveshare_144inch', 'waveshare144inch'): + elif config['ui']['display']['type'] in ( + 'waveshare144lcd', 'ws_144inch', 'ws144inch', 'waveshare_144inch', 'waveshare144inch'): config['ui']['display']['type'] = 'waveshare144lcd' elif config['ui']['display']['type'] in ('ws_213d', 'ws213d', 'waveshare_213d', 'waveshare213d'): @@ -285,7 +286,7 @@ def load_config(args): elif config['ui']['display']['type'] in ('ws_213bc', 'ws213bc', 'waveshare_213bc', 'waveshare213bc'): config['ui']['display']['type'] = 'waveshare213bc' - + elif config['ui']['display']['type'] in ('ws_213bv4', 'ws213bv4', 'waveshare_213bv4', 'waveshare213inb_v4'): config['ui']['display']['type'] = 'waveshare213inb_v4' @@ -318,7 +319,8 @@ def total_unique_handshakes(path): def iface_channels(ifname): channels = [] - output = subprocess.getoutput("/sbin/iwlist %s freq" % ifname) + phy = subprocess.getoutput("/sbin/iw %s info | grep wiphy | cut -d ' ' -f 2" % ifname) + output = subprocess.getoutput("/sbin/iw phy%s channels | grep ' MHz' | sed 's/^.*\[//g' | sed s/\].*\$//g" % phy) for line in output.split("\n"): line = line.strip() if line.startswith("Channel "):