1# Copyright: 2013 Paul Traylor
2# These sources are released under the terms of the MIT license: see LICENSE
3
4import hashlib
5import re
6import time
7
8import gntp.shim
9import gntp.errors as errors
10
11__all__ = [
12	'GNTPRegister',
13	'GNTPNotice',
14	'GNTPSubscribe',
15	'GNTPOK',
16	'GNTPError',
17	'parse_gntp',
18]
19
20#GNTP/<version> <messagetype> <encryptionAlgorithmID>[:<ivValue>][ <keyHashAlgorithmID>:<keyHash>.<salt>]
21GNTP_INFO_LINE = re.compile(
22	'GNTP/(?P<version>\d+\.\d+) (?P<messagetype>REGISTER|NOTIFY|SUBSCRIBE|\-OK|\-ERROR)' +
23	' (?P<encryptionAlgorithmID>[A-Z0-9]+(:(?P<ivValue>[A-F0-9]+))?) ?' +
24	'((?P<keyHashAlgorithmID>[A-Z0-9]+):(?P<keyHash>[A-F0-9]+).(?P<salt>[A-F0-9]+))?\r\n',
25	re.IGNORECASE
26)
27
28GNTP_INFO_LINE_SHORT = re.compile(
29	'GNTP/(?P<version>\d+\.\d+) (?P<messagetype>REGISTER|NOTIFY|SUBSCRIBE|\-OK|\-ERROR)',
30	re.IGNORECASE
31)
32
33GNTP_HEADER = re.compile('([\w-]+):(.+)')
34
35GNTP_EOL = gntp.shim.b('\r\n')
36GNTP_SEP = gntp.shim.b(': ')
37
38
39class _GNTPBuffer(gntp.shim.StringIO):
40	"""GNTP Buffer class"""
41	def writeln(self, value=None):
42		if value:
43			self.write(gntp.shim.b(value))
44		self.write(GNTP_EOL)
45
46	def writeheader(self, key, value):
47		if not isinstance(value, str):
48			value = str(value)
49		self.write(gntp.shim.b(key))
50		self.write(GNTP_SEP)
51		self.write(gntp.shim.b(value))
52		self.write(GNTP_EOL)
53
54
55class _GNTPBase(object):
56	"""Base initilization
57
58	:param string messagetype: GNTP Message type
59	:param string version: GNTP Protocol version
60	:param string encription: Encryption protocol
61	"""
62	def __init__(self, messagetype=None, version='1.0', encryption=None):
63		self.info = {
64			'version': version,
65			'messagetype': messagetype,
66			'encryptionAlgorithmID': encryption
67		}
68		self.hash_algo = {
69			'MD5': hashlib.md5,
70			'SHA1': hashlib.sha1,
71			'SHA256': hashlib.sha256,
72			'SHA512': hashlib.sha512,
73		}
74		self.headers = {}
75		self.resources = {}
76
77	# For Python2 we can just return the bytes as is without worry
78	# but on Python3 we want to make sure we return the packet as
79	# a unicode string so that things like logging won't get confused
80	if gntp.shim.PY2:
81		def __str__(self):
82			return self.encode()
83	else:
84		def __str__(self):
85			return gntp.shim.u(self.encode())
86
87	def _parse_info(self, data):
88		"""Parse the first line of a GNTP message to get security and other info values
89
90		:param string data: GNTP Message
91		:return dict: Parsed GNTP Info line
92		"""
93
94		match = GNTP_INFO_LINE.match(data)
95
96		if not match:
97			raise errors.ParseError('ERROR_PARSING_INFO_LINE')
98
99		info = match.groupdict()
100		if info['encryptionAlgorithmID'] == 'NONE':
101			info['encryptionAlgorithmID'] = None
102
103		return info
104
105	def set_password(self, password, encryptAlgo='MD5'):
106		"""Set a password for a GNTP Message
107
108		:param string password: Null to clear password
109		:param string encryptAlgo: Supports MD5, SHA1, SHA256, SHA512
110		"""
111		if not password:
112			self.info['encryptionAlgorithmID'] = None
113			self.info['keyHashAlgorithm'] = None
114			return
115
116		self.password = gntp.shim.b(password)
117		self.encryptAlgo = encryptAlgo.upper()
118
119		if not self.encryptAlgo in self.hash_algo:
120			raise errors.UnsupportedError('INVALID HASH "%s"' % self.encryptAlgo)
121
122		hashfunction = self.hash_algo.get(self.encryptAlgo)
123
124		password = password.encode('utf8')
125		seed = time.ctime().encode('utf8')
126		salt = hashfunction(seed).hexdigest()
127		saltHash = hashfunction(seed).digest()
128		keyBasis = password + saltHash
129		key = hashfunction(keyBasis).digest()
130		keyHash = hashfunction(key).hexdigest()
131
132		self.info['keyHashAlgorithmID'] = self.encryptAlgo
133		self.info['keyHash'] = keyHash.upper()
134		self.info['salt'] = salt.upper()
135
136	def _decode_hex(self, value):
137		"""Helper function to decode hex string to `proper` hex string
138
139		:param string value: Human readable hex string
140		:return string: Hex string
141		"""
142		result = ''
143		for i in range(0, len(value), 2):
144			tmp = int(value[i:i + 2], 16)
145			result += chr(tmp)
146		return result
147
148	def _decode_binary(self, rawIdentifier, identifier):
149		rawIdentifier += '\r\n\r\n'
150		dataLength = int(identifier['Length'])
151		pointerStart = self.raw.find(rawIdentifier) + len(rawIdentifier)
152		pointerEnd = pointerStart + dataLength
153		data = self.raw[pointerStart:pointerEnd]
154		if not len(data) == dataLength:
155			raise errors.ParseError('INVALID_DATA_LENGTH Expected: %s Recieved %s' % (dataLength, len(data)))
156		return data
157
158	def _validate_password(self, password):
159		"""Validate GNTP Message against stored password"""
160		self.password = password
161		if password is None:
162			raise errors.AuthError('Missing password')
163		keyHash = self.info.get('keyHash', None)
164		if keyHash is None and self.password is None:
165			return True
166		if keyHash is None:
167			raise errors.AuthError('Invalid keyHash')
168		if self.password is None:
169			raise errors.AuthError('Missing password')
170
171		keyHashAlgorithmID = self.info.get('keyHashAlgorithmID','MD5')
172
173		password = self.password.encode('utf8')
174		saltHash = self._decode_hex(self.info['salt'])
175
176		keyBasis = password + saltHash
177		self.key = self.hash_algo[keyHashAlgorithmID](keyBasis).digest()
178		keyHash = self.hash_algo[keyHashAlgorithmID](self.key).hexdigest()
179
180		if not keyHash.upper() == self.info['keyHash'].upper():
181			raise errors.AuthError('Invalid Hash')
182		return True
183
184	def validate(self):
185		"""Verify required headers"""
186		for header in self._requiredHeaders:
187			if not self.headers.get(header, False):
188				raise errors.ParseError('Missing Notification Header: ' + header)
189
190	def _format_info(self):
191		"""Generate info line for GNTP Message
192
193		:return string:
194		"""
195		info = 'GNTP/%s %s' % (
196			self.info.get('version'),
197			self.info.get('messagetype'),
198		)
199		if self.info.get('encryptionAlgorithmID', None):
200			info += ' %s:%s' % (
201				self.info.get('encryptionAlgorithmID'),
202				self.info.get('ivValue'),
203			)
204		else:
205			info += ' NONE'
206
207		if self.info.get('keyHashAlgorithmID', None):
208			info += ' %s:%s.%s' % (
209				self.info.get('keyHashAlgorithmID'),
210				self.info.get('keyHash'),
211				self.info.get('salt')
212			)
213
214		return info
215
216	def _parse_dict(self, data):
217		"""Helper function to parse blocks of GNTP headers into a dictionary
218
219		:param string data:
220		:return dict: Dictionary of parsed GNTP Headers
221		"""
222		d = {}
223		for line in data.split('\r\n'):
224			match = GNTP_HEADER.match(line)
225			if not match:
226				continue
227
228			key = match.group(1).strip()
229			val = match.group(2).strip()
230			d[key] = val
231		return d
232
233	def add_header(self, key, value):
234		self.headers[key] = value
235
236	def add_resource(self, data):
237		"""Add binary resource
238
239		:param string data: Binary Data
240		"""
241		data = gntp.shim.b(data)
242		identifier = hashlib.md5(data).hexdigest()
243		self.resources[identifier] = data
244		return 'x-growl-resource://%s' % identifier
245
246	def decode(self, data, password=None):
247		"""Decode GNTP Message
248
249		:param string data:
250		"""
251		self.password = password
252		self.raw = gntp.shim.u(data)
253		parts = self.raw.split('\r\n\r\n')
254		self.info = self._parse_info(self.raw)
255		self.headers = self._parse_dict(parts[0])
256
257	def encode(self):
258		"""Encode a generic GNTP Message
259
260		:return string: GNTP Message ready to be sent. Returned as a byte string
261		"""
262
263		buff = _GNTPBuffer()
264
265		buff.writeln(self._format_info())
266
267		#Headers
268		for k, v in self.headers.items():
269			buff.writeheader(k, v)
270		buff.writeln()
271
272		#Resources
273		for resource, data in self.resources.items():
274			buff.writeheader('Identifier', resource)
275			buff.writeheader('Length', len(data))
276			buff.writeln()
277			buff.write(data)
278			buff.writeln()
279			buff.writeln()
280
281		return buff.getvalue()
282
283
284class GNTPRegister(_GNTPBase):
285	"""Represents a GNTP Registration Command
286
287	:param string data: (Optional) See decode()
288	:param string password: (Optional) Password to use while encoding/decoding messages
289	"""
290	_requiredHeaders = [
291		'Application-Name',
292		'Notifications-Count'
293	]
294	_requiredNotificationHeaders = ['Notification-Name']
295
296	def __init__(self, data=None, password=None):
297		_GNTPBase.__init__(self, 'REGISTER')
298		self.notifications = []
299
300		if data:
301			self.decode(data, password)
302		else:
303			self.set_password(password)
304			self.add_header('Application-Name', 'pygntp')
305			self.add_header('Notifications-Count', 0)
306
307	def validate(self):
308		'''Validate required headers and validate notification headers'''
309		for header in self._requiredHeaders:
310			if not self.headers.get(header, False):
311				raise errors.ParseError('Missing Registration Header: ' + header)
312		for notice in self.notifications:
313			for header in self._requiredNotificationHeaders:
314				if not notice.get(header, False):
315					raise errors.ParseError('Missing Notification Header: ' + header)
316
317	def decode(self, data, password):
318		"""Decode existing GNTP Registration message
319
320		:param string data: Message to decode
321		"""
322		self.raw = gntp.shim.u(data)
323		parts = self.raw.split('\r\n\r\n')
324		self.info = self._parse_info(self.raw)
325		self._validate_password(password)
326		self.headers = self._parse_dict(parts[0])
327
328		for i, part in enumerate(parts):
329			if i == 0:
330				continue  # Skip Header
331			if part.strip() == '':
332				continue
333			notice = self._parse_dict(part)
334			if notice.get('Notification-Name', False):
335				self.notifications.append(notice)
336			elif notice.get('Identifier', False):
337				notice['Data'] = self._decode_binary(part, notice)
338				#open('register.png','wblol').write(notice['Data'])
339				self.resources[notice.get('Identifier')] = notice
340
341	def add_notification(self, name, enabled=True):
342		"""Add new Notification to Registration message
343
344		:param string name: Notification Name
345		:param boolean enabled: Enable this notification by default
346		"""
347		notice = {}
348		notice['Notification-Name'] = name
349		notice['Notification-Enabled'] = enabled
350
351		self.notifications.append(notice)
352		self.add_header('Notifications-Count', len(self.notifications))
353
354	def encode(self):
355		"""Encode a GNTP Registration Message
356
357		:return string: Encoded GNTP Registration message. Returned as a byte string
358		"""
359
360		buff = _GNTPBuffer()
361
362		buff.writeln(self._format_info())
363
364		#Headers
365		for k, v in self.headers.items():
366			buff.writeheader(k, v)
367		buff.writeln()
368
369		#Notifications
370		if len(self.notifications) > 0:
371			for notice in self.notifications:
372				for k, v in notice.items():
373					buff.writeheader(k, v)
374				buff.writeln()
375
376		#Resources
377		for resource, data in self.resources.items():
378			buff.writeheader('Identifier', resource)
379			buff.writeheader('Length', len(data))
380			buff.writeln()
381			buff.write(data)
382			buff.writeln()
383			buff.writeln()
384
385		return buff.getvalue()
386
387
388class GNTPNotice(_GNTPBase):
389	"""Represents a GNTP Notification Command
390
391	:param string data: (Optional) See decode()
392	:param string app: (Optional) Set Application-Name
393	:param string name: (Optional) Set Notification-Name
394	:param string title: (Optional) Set Notification Title
395	:param string password: (Optional) Password to use while encoding/decoding messages
396	"""
397	_requiredHeaders = [
398		'Application-Name',
399		'Notification-Name',
400		'Notification-Title'
401	]
402
403	def __init__(self, data=None, app=None, name=None, title=None, password=None):
404		_GNTPBase.__init__(self, 'NOTIFY')
405
406		if data:
407			self.decode(data, password)
408		else:
409			self.set_password(password)
410			if app:
411				self.add_header('Application-Name', app)
412			if name:
413				self.add_header('Notification-Name', name)
414			if title:
415				self.add_header('Notification-Title', title)
416
417	def decode(self, data, password):
418		"""Decode existing GNTP Notification message
419
420		:param string data: Message to decode.
421		"""
422		self.raw = gntp.shim.u(data)
423		parts = self.raw.split('\r\n\r\n')
424		self.info = self._parse_info(self.raw)
425		self._validate_password(password)
426		self.headers = self._parse_dict(parts[0])
427
428		for i, part in enumerate(parts):
429			if i == 0:
430				continue  # Skip Header
431			if part.strip() == '':
432				continue
433			notice = self._parse_dict(part)
434			if notice.get('Identifier', False):
435				notice['Data'] = self._decode_binary(part, notice)
436				#open('notice.png','wblol').write(notice['Data'])
437				self.resources[notice.get('Identifier')] = notice
438
439
440class GNTPSubscribe(_GNTPBase):
441	"""Represents a GNTP Subscribe Command
442
443	:param string data: (Optional) See decode()
444	:param string password: (Optional) Password to use while encoding/decoding messages
445	"""
446	_requiredHeaders = [
447		'Subscriber-ID',
448		'Subscriber-Name',
449	]
450
451	def __init__(self, data=None, password=None):
452		_GNTPBase.__init__(self, 'SUBSCRIBE')
453		if data:
454			self.decode(data, password)
455		else:
456			self.set_password(password)
457
458
459class GNTPOK(_GNTPBase):
460	"""Represents a GNTP OK Response
461
462	:param string data: (Optional) See _GNTPResponse.decode()
463	:param string action: (Optional) Set type of action the OK Response is for
464	"""
465	_requiredHeaders = ['Response-Action']
466
467	def __init__(self, data=None, action=None):
468		_GNTPBase.__init__(self, '-OK')
469		if data:
470			self.decode(data)
471		if action:
472			self.add_header('Response-Action', action)
473
474
475class GNTPError(_GNTPBase):
476	"""Represents a GNTP Error response
477
478	:param string data: (Optional) See _GNTPResponse.decode()
479	:param string errorcode: (Optional) Error code
480	:param string errordesc: (Optional) Error Description
481	"""
482	_requiredHeaders = ['Error-Code', 'Error-Description']
483
484	def __init__(self, data=None, errorcode=None, errordesc=None):
485		_GNTPBase.__init__(self, '-ERROR')
486		if data:
487			self.decode(data)
488		if errorcode:
489			self.add_header('Error-Code', errorcode)
490			self.add_header('Error-Description', errordesc)
491
492	def error(self):
493		return (self.headers.get('Error-Code', None),
494			self.headers.get('Error-Description', None))
495
496
497def parse_gntp(data, password=None):
498	"""Attempt to parse a message as a GNTP message
499
500	:param string data: Message to be parsed
501	:param string password: Optional password to be used to verify the message
502	"""
503	data = gntp.shim.u(data)
504	match = GNTP_INFO_LINE_SHORT.match(data)
505	if not match:
506		raise errors.ParseError('INVALID_GNTP_INFO')
507	info = match.groupdict()
508	if info['messagetype'] == 'REGISTER':
509		return GNTPRegister(data, password=password)
510	elif info['messagetype'] == 'NOTIFY':
511		return GNTPNotice(data, password=password)
512	elif info['messagetype'] == 'SUBSCRIBE':
513		return GNTPSubscribe(data, password=password)
514	elif info['messagetype'] == '-OK':
515		return GNTPOK(data)
516	elif info['messagetype'] == '-ERROR':
517		return GNTPError(data)
518	raise errors.ParseError('INVALID_GNTP_MESSAGE')
519