Add Token support for Bugzilla 4.4 and newer

bugz/cli.py

@@ -29,6 +29,7 @@
 """
 
 DEFAULT_COOKIE_FILE = '.bugz_cookie'
+DEFAULT_TOKEN_FILE = '.bugz_token'
 DEFAULT_NUM_COLS = 80
 
 #
@@ -135,6 +136,12 @@ def __init__(self, args):
 		except IOError:
 			pass
 
+		self.token_file = os.path.join(os.environ['HOME'], DEFAULT_TOKEN_FILE)
+		try:
+			self.token = open(self.token_file).read().strip()
+		except IOError:
+			self.token = None
+
 		if getattr(args, 'encoding'):
 			self.enc = args.encoding
 		else:
@@ -151,16 +158,21 @@ def __init__(self, args):
 	def get_input(self, prompt):
 		return raw_input(prompt)
 
+	def set_token(self, *args):
+		if args and self.token:
+			args[0]['Bugzilla_token'] = self.token
+		return args
+
 	def bzcall(self, method, *args):
 		"""Attempt to call method with args. Log in if authentication is required.
 		"""
 		try:
-			return method(*args)
+			return method(*self.set_token(*args))
 		except xmlrpclib.Fault, fault:
 			# Fault code 410 means login required
 			if fault.faultCode == 410 and not self.skip_auth:
 				self.login()
-				return method(*args)
+				return method(*self.set_token(*args))
 			raise
 
 	def login(self, args=None):
@@ -192,10 +204,21 @@ def login(self, args=None):
 			self.bz.User.login(params)
 		except xmlrpclib.Fault as fault:
 			raise BugzError("Can't login: " + fault.faultString)
+		log_info('Logging in')
+		result = self.bz.User.login(params)
+		if 'token' in result:
+			self.token = result['token']
 
 		if args is not None:
-			self.cookiejar.save()
-			os.chmod(self.cookiejar.filename, 0600)
+			if self.token:
+				fd = open(self.token_file, 'w')
+				fd.write(self.token)
+				fd.write('\n')
+				fd.close()
+				os.chmod(self.token_file, 0600)
+			else:
+				self.cookiejar.save()
+				os.chmod(self.cookiejar.filename, 0600)
 
 	def logout(self, args):
 		log_info('logging out')