heat.py

Revision 56, 10.4 kB (checked in by gak, 2 years ago)

Property svn:keywords set to `Id Revision`

Line
1	#!/usr/bin/env python
2	'''
3	wirelessheatmap
4	Copyright 2008 Gerald Kaszuba
5	http://geraldkaszuba.com/
6
7	This program is free software: you can redistribute it and/or modify
8	it under the terms of the GNU General Public License as published by
9	the Free Software Foundation, either version 3 of the License, or
10	(at your option) any later version.
11
12	This program is distributed in the hope that it will be useful,
13	but WITHOUT ANY WARRANTY; without even the implied warranty of
14	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15	GNU General Public License for more details.
16
17	You should have received a copy of the GNU General Public License
18	along with this program. If not, see <http://www.gnu.org/licenses/>.
19	'''
20
21	import colorsys
22	import math
23	import time
24	import pickle
25	from ConfigParser import ConfigParser
26	from optparse import OptionParser
27
28	import pyglet
29	pyglet.options['debug_gl'] = 0
30
31	from pyglet.window import key
32	from pyglet.window import mouse
33	from pyglet import window
34	from pyglet import image
35	from pyglet import font
36	from pyglet.gl import *
37
38	from ap import AccessPoint
39
40	class Heat(object):
41
42	def __init__(self):
43
44	self.parse_config()
45
46	self.window = window.Window(
47	# 640, 480,
48	# 1024, 768,
49	800, 600,
50	# caption='Heat', \
51	# fullscreen=True,
52	)
53	self.window.on_mouse_press = self.on_mouse_press
54	self.window.on_mouse_drag = self.on_mouse_drag
55	self.window.on_mouse_scroll = self.on_mouse_scroll
56	self.window.on_key_release = self.on_key_release
57
58	self.current_ap = ''
59
60	# defaults before image loads
61	self.cellsize = 12
62	self.img_zoom = 1.
63
64	if self.cfg.image_file:
65	self.img = image.load(self.cfg.image_file)
66	self.img_zoom = float(self.window.width) / self.img.width
67	zoomier = float(self.window.height) / self.img.height
68	if zoomier < self.img_zoom:
69	self.img_zoom = zoomier
70
71	self.cellsize = self.img.width / 100
72
73	else:
74	self.img = None
75
76	self.img_pos = [0., 0.]
77
78	self.font_size = 20
79	self.font = font.load('Arial', self.font_size)
80
81	self.load_data()
82
83	self.last_dump_read = 0
84
85	self.samples = 10
86	self.grid = {}
87
88	self.grid_margin = 100 # pixels surrounding the image
89
90	def parse_config(self):
91	self.cfg = ConfigParser()
92	self.opt = OptionParser(usage='usage: %prog [options]')
93
94	# Configuration loading and saving
95	self.opt.add_option('--config', dest='config_file', \
96	help='Read configuration from a file.')
97	self.opt.add_option('--save', dest='dest_cfg', \
98	help='Generate a config file from your options into a file.')
99
100	# General shit
101	self.opt.add_option('--image', dest='image_file')
102	self.opt.add_option('--dump', dest='dump_file')
103	self.opt.add_option('--store', dest='store_file')
104
105	(opts, args) = self.opt.parse_args()
106
107	# If we want to load from a cfg file, update the defaults then override
108	# them with the command line arguments.
109	if opts.config_file:
110	self.cfg.read(opts.config_file)
111	for k, v in self.cfg.items('main'):
112	# Since these are saved as strings we have to cast them back.
113	# XXX: This needs to be cleaner...
114	if v == 'False':
115	v = False
116	if v == 'True':
117	v = True
118	self.opt.set_default(k, v)
119	(opts, args) = self.opt.parse_args()
120
121	if opts.dest_cfg:
122	f = open(opts.dest_cfg, 'w')
123	f.write('[main]\n')
124	for a in self.opt.defaults:
125	if a in ('dest_cfg', 'config_file'):
126	continue
127	f.write('%s = %s\n' % (a, getattr(opts, a)))
128	f.close()
129	print '%s written' % opts.dest_cfg
130	sys.exit(0)
131
132	self.cfg = opts
133
134	return args
135
136	def draw_text(self, text, line):
137	t = font.Text(self.font, text, y=line*self.font_size)
138	t.draw()
139
140	def read_dump(self):
141	if not self.cfg.dump_file:
142	return
143	d = open(self.cfg.dump_file).read()
144	lines = d.split('\n')
145
146	for line in lines:
147
148	bits = line.split(',')
149
150	if len(bits) != 15:
151	continue
152	if bits[0] == 'BSSID':
153	continue
154
155	ap = AccessPoint(*bits)
156
157	self.ap[ap.bssid] = ap
158	if ap.bssid not in self.store:
159	self.store[ap.bssid] = {}
160
161	# Make the first read AP the currently selected one
162	if not self.current_ap:
163	self.current_ap = ap.bssid
164
165	def go(self):
166	while not self.window.has_exit:
167
168	if self.last_dump_read + 1 < time.time():
169	self.read_dump()
170	self.last_dump_read = time.time()
171
172	self.window.clear()
173
174	if self.current_ap:
175	self.calc_heat()
176
177	self.map_matrix()
178	self.draw_heat()
179	self.draw_strength()
180
181	glLoadIdentity()
182	self.draw_text('%s %s %f' % (self.current_ap,
183	self.ap[self.current_ap].essid,
184	self.ap[self.current_ap].power), 0)
185
186	if self.img:
187	self.map_matrix()
188	glEnable(GL_BLEND)
189	glBlendFunc(GL_SRC_ALPHA, GL_ONE)
190	glColor4f(1, 1, 1, 0.5)
191	self.img.blit(0, 0)
192	glDisable(GL_BLEND)
193
194	self.window.flip()
195	self.window.dispatch_events()
196	# time.sleep(0.5)
197
198	def map_matrix(self):
199	glLoadIdentity()
200	glTranslatef(self.img_pos[0], self.img_pos[1], 0)
201	glScalef(self.img_zoom, self.img_zoom, 1)
202
203	def screen2map(self, x, y):
204	x -= self.img_pos[0]
205	y -= self.img_pos[1]
206	x /= self.img_zoom
207	y /= self.img_zoom
208	return x, y
209
210	def draw_strength(self):
211	if not self.current_ap:
212	return
213	if self.current_ap not in self.store:
214	return
215	if not self.store[self.current_ap]:
216	return
217
218	max_str = max(self.store[self.current_ap].values())
219	min_str = min(self.store[self.current_ap].values())
220
221	if not max_str:
222	return
223	if max_str == min_str:
224	return
225
226	glPointSize(2)
227	glBegin(GL_POINTS)
228	for pos, strength in self.store[self.current_ap].items():
229	s = 0.5 + (strength - min_str) / (max_str - min_str)
230	glColor3f(s, s, s)
231	glVertex2f(pos[0], pos[1], 0)
232	glEnd()
233
234	def calc_heat(self):
235
236	if self.grid:
237	return
238
239	ap = self.store[self.current_ap]
240	self.grid = {}
241
242	gmin = -self.grid_margin
243	gmaxx = self.img.width + self.grid_margin
244	gmaxy = self.img.height + self.grid_margin
245
246	for x in range(gmin, gmaxx, self.cellsize):
247	for y in range(gmin, gmaxy, self.cellsize):
248
249	distance_strengths = []
250	for pos, strength in ap.items():
251	dist = math.hypot(x - pos[0], y - pos[1])
252	distance_strengths.append((dist, strength))
253
254	heat = 0
255	def srt(a, b):
256	if a[0] > b[0]:
257	return 1
258	return -1
259
260	distance_strengths.sort(srt)
261
262	for d, s in distance_strengths[:self.samples]:
263	heat += s
264	heat /= self.samples
265	self.grid[x, y] = heat
266
267	def draw_heat(self):
268	if not self.grid:
269	return
270	max_heat = max(self.grid.values())
271	min_heat = min(self.grid.values())
272	if not max_heat:
273	return
274	if max_heat == min_heat:
275	return
276
277	glBegin(GL_QUADS)
278
279	for pos, heat in self.grid.items():
280	x = pos[0] - self.cellsize / 2
281	y = pos[1] - self.cellsize / 2
282	heat -= min_heat
283	heat /= (max_heat - min_heat)
284	rgb = self.col(heat)
285	if not rgb:
286	continue
287
288	glColor3f(*rgb)
289	glVertex2f(x, y)
290	glVertex2f(x + self.cellsize, y)
291	glVertex2f(x + self.cellsize, y + self.cellsize)
292	glVertex2f(x, y + self.cellsize)
293
294	glEnd()
295
296	def col(self, val):
297	return colorsys.hsv_to_rgb((1 - val) * 2. / 3, 1, val * 0.6)
298
299	def on_mouse_press(self, x, y, button, modifiers):
300	if button == mouse.LEFT:
301	x, y = self.screen2map(x, y)
302	for ap in self.ap.values():
303	self.store[ap.bssid][x, y] = ap.power
304	self.save_data()
305
306	def on_mouse_drag(self, x, y, dx, dy, button, modifiers):
307	if button == mouse.RIGHT:
308	self.img_pos[0] += dx
309	self.img_pos[1] += dy
310
311	def on_mouse_scroll(self, x, y, scroll_x, scroll_y):
312	if scroll_y > 0:
313	self.img_zoom /= 0.8 * scroll_y
314	if scroll_y < 0:
315	self.img_zoom = 0.8 -scroll_y
316
317	def on_key_release(self, symbol, mods):
318	self.grid = {}
319	if symbol == key.S:
320	pyglet.image.get_buffer_manager().get_color_buffer(). \
321	save('screenshot.png')
322	print 'Saved as screenshot.png'
323	if symbol == key.SPACE:
324
325	bssids = self.ap.keys()
326	bssids.sort()
327
328	try:
329	pos = bssids.index(self.current_ap)
330	except ValueError:
331	pos = 0
332
333	try:
334	self.current_ap = bssids[pos + 1]
335	except IndexError:
336	self.current_ap = bssids[0]
337
338	if symbol == key.UP:
339	self.samples += 1
340	print self.samples
341	if symbol == key.DOWN:
342	self.samples -= 1
343	print self.samples
344
345	def load_data(self):
346	# store[essid][x, y] = strength
347	# ap[essid] = AccessPoint()
348	try:
349	self.store, self.ap = pickle.load(open(self.cfg.store_file))
350	except Exception:
351	self.store = {}
352	self.ap = {}
353	print 'Creating new store'
354	return
355
356	print self.cfg.store_file, 'loaded'
357	print len(self.store), 'stores'
358	print self.ap
359
360	def save_data(self):
361	pickle.dump((self.store, self.ap), open(self.cfg.store_file, 'wb'))
362
363
364	Heat().go()

Note: See TracBrowser for help on using the browser.