<commit_before>import sys
import DiffEngine
import os
import hashlib
import subprocess
import dircache
import shutil
from _winreg import *
import tempfile
import win32api
import win32security
import win32con

LogToStdout = 0x1
LogToDbgview = 0x2
LogToFile = 0x4
LogToIDAMessageBox = 0x8

LOG_DARUNGRIM=0x00000001
LOG_DIFF_MACHINE=0x00000002
LOG_IDA_CONTROLLER=0x00000004
LOG_SQL=0x00000008
LOG_ONE_LOCATION_INFO=0x0000000F
LOG_MATCH_RATE=0x00000010

class DarunGrim:
	DebugLevel=0
	def __init__ ( self, src_filename='', target_filename='', ida_path='', start_ida_listener=False):
		self.SrcStorage=''
		self.TargetStorage=''
		self.IDAPath=''
		self.DarunGrim = DiffEngine.DarunGrim()
		self.DarunGrim.SetLogParameters(LogToStdout, 10, "")

		self.DarunGrim.AddSrcDumpAddress(0)
		self.DarunGrim.AddTargetDumpAddress(0)
		self.DarunGrim.EnableLogType(LOG_DIFF_MACHINE)

		if src_filename:
			self.SetSourceFilename(src_filename)

		if target_filename:
			self.SetTargetFilename(target_filename)

		self.SetIDAPath()
		self.ListeningPort=0

		if start_ida_listener:
			self.ListeningPort=self.DarunGrim.StartIDAListenerThread(self.ListeningPort)
		self.DGFSotrage=''

	def SetIDAPath(self,ida_path='',is_64=False):
		if ida_path=='' or not os.path.isfile(ida_path):
			if self.DebugLevel>0:
				print 'Locating IDA executables...'

			ida_executables=self.LocateIDAExecutables()
			if len(ida_executables)>0:
				ida_path=ida_executables[0][0]

				if self.DebugLevel>0:
					print 'Using IDA executable [%s] ...' % ida_path

		if ida_path:
			ida_path=str(ida_path)
			if is_64:
				self.IDA64Path=ida_path
			else:
				self.IDAPath=ida_path

			self.DarunGrim.SetIDAPath(ida_path,is_64)
		else:
			if self.DebugLevel>0:
				print 'No IDA found'

	def CheckIDAPlugin(self):
		darungrim_plugin_path=os.path.join(os.path.dirname(self.IDAPath),"plugins/DarunGrimPlugin.plw")
		if os.path.isfile(darungrim_plugin_path):
			return True
		return False

	def InstallIDAPlugin(self,filename):
		if not self.IDAPath:
			return (False, 'Invalid IDA path')

		darungrim_plugin_path=os.path.join(os.path.dirname(self.IDAPath),"plugins\\DarunGrimPlugin.plw")

		ph = win32api.GetCurrentProcess()
		th = win32security.OpenProcessToken(ph,win32con.MAXIMUM_ALLOWED)

		virtual_store_enabled=win32security.GetTokenInformation(th, win32security.TokenVirtualizationEnabled)
		if virtual_store_enabled:
			return (False,'Program files are under Virtual Store.')

		try:
 			shutil.copy(filename, darungrim_plugin_path)
		except Exception as e:
			s = str(e)
			return (False, s)

		if not os.path.isfile(darungrim_plugin_path):
			return (False,'Not copied')

		return (True,"File location: %s" % darungrim_plugin_path)

	def LocateIDAExecutables(self,is_64=False):
		if is_64:
			executables=['idaq64.exe','idag64.exe']
		else:
			executables=['idaq.exe','idag.exe']

		ida_executables=[]
		for (ida_path,ctime) in self.LocateIDAInstallations():
			for executable in executables:
				filename=os.path.join(ida_path,executable)
				if os.path.isfile(filename):
					ida_executables.append([filename, os.path.getctime(filename)])

		ida_executables=sorted(ida_executables, key=lambda x:x[1], reverse=True)

		return ida_executables

	def LocateIDAInstallations(self):
		ida_paths={}
		programs_files={}
		programs_files[os.environ['ProgramFiles(x86)']]=1
		programs_files[os.environ['ProgramFiles']]=1

		for programs_file in programs_files.keys():
			for dir in dircache.listdir(programs_file):
				if dir[0:4]=='IDA ':
					ida_paths[os.path.join(programs_file,dir)]=1

		reg=ConnectRegistry(None,HKEY_LOCAL_MACHINE)
		key=OpenKey(reg,r"SOFTWARE\Wow6432Node\Microsoft\Windows\CurrentVersion\Uninstall")

		i=0
		while True:
			try:
				subkey=EnumKey(key,i)
				if subkey[0:4]=='IDA ':
					ida_key=OpenKey(key,subkey)
					(value,type)=QueryValueEx(ida_key,"Inno Setup: App Path")
					ida_paths[value]=1
				i+=1
			except:
				break
		"""
		TODO:
			HKEY_CLASSES_ROOT\Applications\idaq.exe\shell\open\command
			HKEY_CLASSES_ROOT\IDApro.Database32\shell\open\command
			HKEY_CLASSES_ROOT\IDApro.Database64\shell\open\command
			HKEY_CLASSES_ROOT\Local Settings\Software\Microsoft\Windows\Shell\MuiCache
			HKEY_CURRENT_USER\Software\Hex-Rays\IDA
			HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\.i64\OpenWithList
			HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\.idb\OpenWithList
			HKEY_CURRENT_USER\Software\Hex-Rays
		"""

		folders=[]
		for folder in ida_paths.keys():
			folders.append([folder, os.path.getctime(folder)])

		folders=sorted(folders, key=lambda x:x[1], reverse=True)

		return folders

	def InstallIDAPlugins(self,plugin_filename):
		for [ida_path, plugin_path] in self.CheckIDAPlugins():
			if self.DebugLevel>0:
				print "Installing DarunGrim Plugin %s -> %s" % (plugin_filename, plugin_path)
			shutil.copy(plugin_filename, plugin_path)

	def CheckIDAPlugins(self):
		missing_plugins=[]
		for (ida_path,ctime) in self.LocateIDAInstallations():
			plugin_path=os.path.join(os.path.join(ida_path,"plugins"),"DarunGrimPlugin.plw")
			if not os.path.isfile(plugin_path):
				if self.DebugLevel>0:
					print 'DarunGrim Plugin missing: ', plugin_path
				missing_plugins.append([ida_path, plugin_path])

		return missing_plugins

	def SetSourceController(self,identity):
		self.DarunGrim.SetSourceController(str(identity))

	def SetTargetController(self,identity):
		self.DarunGrim.SetTargetController(str(identity))

	def SetSourceFilename(self,src_filename, is_dgf=False):
		self.SrcFilename = str(src_filename)

		if self.SrcFilename:
			self.DarunGrim.SetSourceFilename( self.SrcFilename )

	def SetTargetFilename(self,target_filename, is_dgf=False):
		self.TargetFilename = str(target_filename)

		if self.TargetFilename:
			self.DarunGrim.SetTargetFilename( self.TargetFilename )

	def SetLogFile(self,log_filename,log_level=10):
		self.DarunGrim.SetLogParameters(LogToFile, log_level, str(log_filename))

	def SetDGFSotrage(self,dgf_dir):
		self.DGFSotrage=dgf_dir

	def GetDGFName(self,filename):
		if self.DGFSotrage:
			fd=open(filename,'rb')
			data=fd.read()
			fd.close()

			s=hashlib.sha1()
			s.update(data)

			filename=os.path.join(self.DGFSotrage,"%s.dgf" % s.hexdigest())
		else:
			if filename[-4]=='.':
				filename=filename[0:-4] + ".dgf"
			else:
				filename=filename + '.dgf'

		return filename

	def Is64(self,filename):
		import pefile
		pe = pefile.PE(filename)
		_32bitFlag = pefile.IMAGE_CHARACTERISTICS['IMAGE_FILE_32BIT_MACHINE']

		if ( pe.FILE_HEADER.Machine & _32bitFlag ) == _32bitFlag:
			return False
		else:
			return True

	def SetStorageNames(self, src_storage, target_storage):
		self.SrcStorage=str(src_storage)
		self.TargetStorage=str(target_storage)

	def PerformDiff( self, output_storage, src_ida_log_filename = "src.log", target_ida_log_filename = "target.log" ):
		if not self.SrcStorage:
			self.SrcStorage=self.GetDGFName(self.SrcFilename)
			src_ida_log_filename=os.path.join( os.getcwd(), str(src_ida_log_filename))

			if not os.path.isfile(self.SrcStorage) or os.path.getsize(self.SrcStorage)==0:
				if self.Is64(self.SrcFilename):
					src_is_64=True
				else:
					src_is_64=False

				print 'Generate DarunGrim data file %s -> %s' % (self.SrcFilename, self.SrcStorage)
				self.DarunGrim.GenerateSourceDGFFromIDA(self.SrcStorage, src_ida_log_filename, src_is_64)

		if not self.TargetStorage:
			self.TargetStorage=self.GetDGFName(self.TargetFilename)

			output_storage=str(output_storage)

			target_ida_log_filename=os.path.join( os.getcwd(), str(target_ida_log_filename))

			if not os.path.isfile(self.TargetStorage) or os.path.getsize(self.TargetStorage)==0:
				if self.Is64(self.TargetFilename):
					target_is_64=True
				else:
					target_is_64=False

				print 'Generate DarunGrim data file %s -> %s' % (self.TargetFilename, self.TargetStorage)
				self.DarunGrim.GenerateTargetDGFFromIDA(self.TargetStorage, target_ida_log_filename, target_is_64)

		print 'Perform Diffing...'
		self.DarunGrim.PerformDiff(str(self.SrcStorage), 0, str(self.TargetStorage), 0, str(output_storage))

	def OpenIDA(self,filename):
		if filename[-4:].lower()=='.idb':
			is_64=False
		elif filename[-4:].lower()=='.i64':
			is_64=True
		else:
			if self.Is64(filename):
				is_64=True
			else:
				is_64=False

		if is_64:
			ida_path=self.IDA64Path
		else:
			ida_path=self.IDAPath

		fd=tempfile.TemporaryFile(delete=False)

		idc_data="static main()\n" + \
				"{\n" + \
				"	Wait();\n" + \
				"	RunPlugin( \"DarunGrimPlugin\", 1 );\n" + \
				"	ConnectToDarunGrim(%d);\n" + \
				"}"

		fd.write(idc_data % self.ListeningPort)
		idc_filename=fd.name
		fd.close()

		subprocess.Popen([ida_path, "-S" + idc_filename, filename])

	def SyncIDA( self ):
		self.DarunGrim.AcceptIDAClientsFromSocket()

	def JumpToAddresses( self, source_address, target_address ):
		self.DarunGrim.JumpToAddresses( source_address, target_address )

	def ColorAddress( self, type, start_address, end_address, color ):
		self.DarunGrim.ColorAddress( type, start_address, end_address, color )

if __name__ == '__main__':
	from optparse import OptionParser
	import pprint

	parser = OptionParser()

	parser.add_option('-d','--diff',
					dest='diff',help="Perform diff",
					action="store_true", default=False,
					metavar="DIFF")

	parser.add_option('-i','--ida',
					dest='ida',help="Locate IDA",
					action="store_true", default=False,
					metavar="IDA")

	parser.add_option("-s", "--source_address", dest="source_address",
						help="Source function address", type="int", default=0, metavar="SOURCE_ADDRESS")

	parser.add_option("-t", "--target_address", dest="target_address",
						help="Target function address", type="int", default=0, metavar="TARGET_ADDRESS")

	(options, args) = parser.parse_args()

	if options.diff:
		src_storage = args[0]
		target_storage = args[1]
		result_storage = args[2]

		darungrim=DarunGrim(src_storage, target_storage)
		darungrim.SetDGFSotrage(os.getcwd())
		darungrim.PerformDiff("out.dgf")

	elif options.ida:
		darungrim=DarunGrim()
		pprint.pprint(darungrim.LocateIDAInstallations())
		darungrim.InstallIDAPlugin(r'C:\mat\Src\DarunGrim\Release\DarunGrimPlugin.plw')
		ida_executables=darungrim.LocateIDAExecutables()

		pprint.pprint(ida_executables)
<commit_msg>Wrapping with str for swig arguments<commit_after>import sys
import DiffEngine
import os
import hashlib
import subprocess
import dircache
import shutil
from _winreg import *
import tempfile
import win32api
import win32security
import win32con

LogToStdout = 0x1
LogToDbgview = 0x2
LogToFile = 0x4
LogToIDAMessageBox = 0x8

LOG_DARUNGRIM=0x00000001
LOG_DIFF_MACHINE=0x00000002
LOG_IDA_CONTROLLER=0x00000004
LOG_SQL=0x00000008
LOG_ONE_LOCATION_INFO=0x0000000F
LOG_MATCH_RATE=0x00000010

class DarunGrim:
	DebugLevel=0
	def __init__ ( self, src_filename='', target_filename='', ida_path='', start_ida_listener=False):
		self.SrcStorage=''
		self.TargetStorage=''
		self.IDAPath=''
		self.DarunGrim = DiffEngine.DarunGrim()
		self.DarunGrim.SetLogParameters(LogToStdout, 10, "")

		self.DarunGrim.AddSrcDumpAddress(0)
		self.DarunGrim.AddTargetDumpAddress(0)
		self.DarunGrim.EnableLogType(LOG_DIFF_MACHINE)

		if src_filename:
			self.SetSourceFilename(src_filename)

		if target_filename:
			self.SetTargetFilename(target_filename)

		self.SetIDAPath()
		self.ListeningPort=0

		if start_ida_listener:
			self.ListeningPort=self.DarunGrim.StartIDAListenerThread(self.ListeningPort)
		self.DGFSotrage=''

	def SetIDAPath(self,ida_path='',is_64=False):
		if ida_path=='' or not os.path.isfile(ida_path):
			if self.DebugLevel>0:
				print 'Locating IDA executables...'

			ida_executables=self.LocateIDAExecutables()
			if len(ida_executables)>0:
				ida_path=ida_executables[0][0]

				if self.DebugLevel>0:
					print 'Using IDA executable [%s] ...' % ida_path

		if ida_path:
			ida_path=str(ida_path)
			if is_64:
				self.IDA64Path=ida_path
			else:
				self.IDAPath=ida_path

			self.DarunGrim.SetIDAPath(ida_path,is_64)
		else:
			if self.DebugLevel>0:
				print 'No IDA found'

	def CheckIDAPlugin(self):
		darungrim_plugin_path=os.path.join(os.path.dirname(self.IDAPath),"plugins/DarunGrimPlugin.plw")
		if os.path.isfile(darungrim_plugin_path):
			return True
		return False

	def InstallIDAPlugin(self,filename):
		if not self.IDAPath:
			return (False, 'Invalid IDA path')

		darungrim_plugin_path=os.path.join(os.path.dirname(self.IDAPath),"plugins\\DarunGrimPlugin.plw")

		ph = win32api.GetCurrentProcess()
		th = win32security.OpenProcessToken(ph,win32con.MAXIMUM_ALLOWED)

		virtual_store_enabled=win32security.GetTokenInformation(th, win32security.TokenVirtualizationEnabled)
		if virtual_store_enabled:
			return (False,'Program files are under Virtual Store.')

		try:
 			shutil.copy(filename, darungrim_plugin_path)
		except Exception as e:
			s = str(e)
			return (False, s)

		if not os.path.isfile(darungrim_plugin_path):
			return (False,'Not copied')

		return (True,"File location: %s" % darungrim_plugin_path)

	def LocateIDAExecutables(self,is_64=False):
		if is_64:
			executables=['idaq64.exe','idag64.exe']
		else:
			executables=['idaq.exe','idag.exe']

		ida_executables=[]
		for (ida_path,ctime) in self.LocateIDAInstallations():
			for executable in executables:
				filename=os.path.join(ida_path,executable)
				if os.path.isfile(filename):
					ida_executables.append([filename, os.path.getctime(filename)])

		ida_executables=sorted(ida_executables, key=lambda x:x[1], reverse=True)

		return ida_executables

	def LocateIDAInstallations(self):
		ida_paths={}
		programs_files={}
		programs_files[os.environ['ProgramFiles(x86)']]=1
		programs_files[os.environ['ProgramFiles']]=1

		for programs_file in programs_files.keys():
			for dir in dircache.listdir(programs_file):
				if dir[0:4]=='IDA ':
					ida_paths[os.path.join(programs_file,dir)]=1

		reg=ConnectRegistry(None,HKEY_LOCAL_MACHINE)
		key=OpenKey(reg,r"SOFTWARE\Wow6432Node\Microsoft\Windows\CurrentVersion\Uninstall")

		i=0
		while True:
			try:
				subkey=EnumKey(key,i)
				if subkey[0:4]=='IDA ':
					ida_key=OpenKey(key,subkey)
					(value,type)=QueryValueEx(ida_key,"Inno Setup: App Path")
					ida_paths[value]=1
				i+=1
			except:
				break
		"""
		TODO:
			HKEY_CLASSES_ROOT\Applications\idaq.exe\shell\open\command
			HKEY_CLASSES_ROOT\IDApro.Database32\shell\open\command
			HKEY_CLASSES_ROOT\IDApro.Database64\shell\open\command
			HKEY_CLASSES_ROOT\Local Settings\Software\Microsoft\Windows\Shell\MuiCache
			HKEY_CURRENT_USER\Software\Hex-Rays\IDA
			HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\.i64\OpenWithList
			HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\.idb\OpenWithList
			HKEY_CURRENT_USER\Software\Hex-Rays
		"""

		folders=[]
		for folder in ida_paths.keys():
			folders.append([folder, os.path.getctime(folder)])

		folders=sorted(folders, key=lambda x:x[1], reverse=True)

		return folders

	def InstallIDAPlugins(self,plugin_filename):
		for [ida_path, plugin_path] in self.CheckIDAPlugins():
			if self.DebugLevel>0:
				print "Installing DarunGrim Plugin %s -> %s" % (plugin_filename, plugin_path)
			shutil.copy(plugin_filename, plugin_path)

	def CheckIDAPlugins(self):
		missing_plugins=[]
		for (ida_path,ctime) in self.LocateIDAInstallations():
			plugin_path=os.path.join(os.path.join(ida_path,"plugins"),"DarunGrimPlugin.plw")
			if not os.path.isfile(plugin_path):
				if self.DebugLevel>0:
					print 'DarunGrim Plugin missing: ', plugin_path
				missing_plugins.append([ida_path, plugin_path])

		return missing_plugins

	def SetSourceController(self,identity):
		self.DarunGrim.SetSourceController(str(identity))

	def SetTargetController(self,identity):
		self.DarunGrim.SetTargetController(str(identity))

	def SetSourceFilename(self,src_filename, is_dgf=False):
		self.SrcFilename = str(src_filename)

		if self.SrcFilename:
			self.DarunGrim.SetSourceFilename( self.SrcFilename )

	def SetTargetFilename(self,target_filename, is_dgf=False):
		self.TargetFilename = str(target_filename)

		if self.TargetFilename:
			self.DarunGrim.SetTargetFilename( self.TargetFilename )

	def SetLogFile(self,log_filename,log_level=10):
		self.DarunGrim.SetLogParameters(LogToFile, log_level, str(log_filename))

	def SetDGFSotrage(self,dgf_dir):
		self.DGFSotrage=dgf_dir

	def GetDGFName(self,filename):
		if self.DGFSotrage:
			fd=open(filename,'rb')
			data=fd.read()
			fd.close()

			s=hashlib.sha1()
			s.update(data)

			filename=os.path.join(self.DGFSotrage,"%s.dgf" % s.hexdigest())
		else:
			if filename[-4]=='.':
				filename=filename[0:-4] + ".dgf"
			else:
				filename=filename + '.dgf'

		return filename

	def Is64(self,filename):
		import pefile
		pe = pefile.PE(filename)
		_32bitFlag = pefile.IMAGE_CHARACTERISTICS['IMAGE_FILE_32BIT_MACHINE']

		if ( pe.FILE_HEADER.Machine & _32bitFlag ) == _32bitFlag:
			return False
		else:
			return True

	def SetStorageNames(self, src_storage, target_storage):
		self.SrcStorage=str(src_storage)
		self.TargetStorage=str(target_storage)

	def PerformDiff( self, output_storage, src_ida_log_filename = "src.log", target_ida_log_filename = "target.log" ):
		if not self.SrcStorage:
			self.SrcStorage=self.GetDGFName(self.SrcFilename)
			src_ida_log_filename=os.path.join( os.getcwd(), str(src_ida_log_filename))

			if not os.path.isfile(self.SrcStorage) or os.path.getsize(self.SrcStorage)==0:
				if self.Is64(self.SrcFilename):
					src_is_64=True
				else:
					src_is_64=False

				print 'Generate DarunGrim data file %s -> %s' % (self.SrcFilename, self.SrcStorage)
				self.DarunGrim.GenerateSourceDGFFromIDA(str(self.SrcStorage), str(src_ida_log_filename), src_is_64)

		if not self.TargetStorage:
			self.TargetStorage=self.GetDGFName(self.TargetFilename)

			output_storage=str(output_storage)

			target_ida_log_filename=os.path.join( os.getcwd(), str(target_ida_log_filename))

			if not os.path.isfile(self.TargetStorage) or os.path.getsize(self.TargetStorage)==0:
				if self.Is64(self.TargetFilename):
					target_is_64=True
				else:
					target_is_64=False

				print 'Generate DarunGrim data file %s -> %s' % (self.TargetFilename, self.TargetStorage)
				self.DarunGrim.GenerateTargetDGFFromIDA(str(self.TargetStorage), str(target_ida_log_filename), target_is_64)

		print 'Perform Diffing...'
		self.DarunGrim.PerformDiff(str(self.SrcStorage), 0, str(self.TargetStorage), 0, str(output_storage))

	def OpenIDA(self,filename):
		if filename[-4:].lower()=='.idb':
			is_64=False
		elif filename[-4:].lower()=='.i64':
			is_64=True
		else:
			if self.Is64(filename):
				is_64=True
			else:
				is_64=False

		if is_64:
			ida_path=self.IDA64Path
		else:
			ida_path=self.IDAPath

		fd=tempfile.TemporaryFile(delete=False)

		idc_data="static main()\n" + \
				"{\n" + \
				"	Wait();\n" + \
				"	RunPlugin( \"DarunGrimPlugin\", 1 );\n" + \
				"	ConnectToDarunGrim(%d);\n" + \
				"}"

		fd.write(idc_data % self.ListeningPort)
		idc_filename=fd.name
		fd.close()

		subprocess.Popen([ida_path, "-S" + idc_filename, filename])

	def SyncIDA( self ):
		self.DarunGrim.AcceptIDAClientsFromSocket()

	def JumpToAddresses( self, source_address, target_address ):
		self.DarunGrim.JumpToAddresses( source_address, target_address )

	def ColorAddress( self, type, start_address, end_address, color ):
		self.DarunGrim.ColorAddress( type, start_address, end_address, color )

if __name__ == '__main__':
	from optparse import OptionParser
	import pprint

	parser = OptionParser()

	parser.add_option('-d','--diff',
					dest='diff',help="Perform diff",
					action="store_true", default=False,
					metavar="DIFF")

	parser.add_option('-i','--ida',
					dest='ida',help="Locate IDA",
					action="store_true", default=False,
					metavar="IDA")

	parser.add_option("-s", "--source_address", dest="source_address",
						help="Source function address", type="int", default=0, metavar="SOURCE_ADDRESS")

	parser.add_option("-t", "--target_address", dest="target_address",
						help="Target function address", type="int", default=0, metavar="TARGET_ADDRESS")

	(options, args) = parser.parse_args()

	if options.diff:
		src_storage = args[0]
		target_storage = args[1]
		result_storage = args[2]

		darungrim=DarunGrim(src_storage, target_storage)
		darungrim.SetDGFSotrage(os.getcwd())
		darungrim.PerformDiff("out.dgf")

	elif options.ida:
		darungrim=DarunGrim()
		pprint.pprint(darungrim.LocateIDAInstallations())
		darungrim.InstallIDAPlugin(r'C:\mat\Src\DarunGrim\Release\DarunGrimPlugin.plw')
		ida_executables=darungrim.LocateIDAExecutables()

		pprint.pprint(ida_executables)
<|endoftext|>


<commit_before>from modules import *
channel = ['#nurfed_x', "#briarcliff"]

class Core(Module):
    """Core Module to provide chat commands for core functionality of
    ppbot.

    """

    def __init__(self, server):
        """Constructor."""

        Module.__init__(self, server)
        self.module_handler = None

    def _register_events(self):
        """Register module commands."""

        # events
        self.add_event('welcome', 'welcome')

        # commands
        self.add_command('load')
        self.add_command('reload')
        self.add_command('raw')

    @master
    def load(self, event):
        """Loads a module specified from IRC.

        @param the irclib event object

        """

        if self.num_args == 1:
            if self.module_handler.load(event['args'][0].title()):
                self.notice(event['nick'],
                    'Successfully loaded module: %s' % event['args'][0])
            else:
                self.notice(event['nick'],
                    'Could not load module: %s' % event['args'][0])
        else:
            self.syntax_message(event['nick'], '.load <module>')

    @master
    def reload(self, event):
        """Reloads a module specified from IRC.

        @param the irclib event object

        """

        if self.num_args == 1:
            # if module was successfully reloaded
            if self.module_handler.reload(event['args'][0].title()):
                # alert the user that reload was successful
                self.notice(event['nick'],
                    'Successfully reloaded module: %s' % event['args'][0])
                # re-register any commands
                self.module_handler.modules[event['args'][0].title()]._register_events()
            else:
                # alert user that reload failed
                self.notice(event['nick'],
                    'Could not reload module: %s' % event['args'][0])
        else:

            self.syntax_message(event['nick'], '.reload <module>')

    @owner
    def raw(self, event):
        """Processes a raw python command from IRC.

        @param the python command to run

        """
        try:
            print ' '.join(event['args'])
            self.msg(event['target'], eval(' '.join(event['args'])))
        except Exception, e:
            print e

    def welcome(self, event):
        """Event handler for when the bot connects to a server."""

        print "Connected to network: %s." % self.server.server_config.name

    	for channel in self.server.server_config.channels:
    	    self.server.join(channel.name)

<commit_msg>Removing old channel list<commit_after>from modules import *

class Core(Module):
    """Core Module to provide chat commands for core functionality of
    ppbot.

    """

    def __init__(self, server):
        """Constructor."""

        Module.__init__(self, server)
        self.module_handler = None

    def _register_events(self):
        """Register module commands."""

        # events
        self.add_event('welcome', 'welcome')

        # commands
        self.add_command('load')
        self.add_command('reload')
        self.add_command('raw')

    @master
    def load(self, event):
        """Loads a module specified from IRC.

        @param the irclib event object

        """

        if self.num_args == 1:
            if self.module_handler.load(event['args'][0].title()):
                self.notice(event['nick'],
                    'Successfully loaded module: %s' % event['args'][0])
            else:
                self.notice(event['nick'],
                    'Could not load module: %s' % event['args'][0])
        else:
            self.syntax_message(event['nick'], '.load <module>')

    @master
    def reload(self, event):
        """Reloads a module specified from IRC.

        @param the irclib event object

        """

        if self.num_args == 1:
            # if module was successfully reloaded
            if self.module_handler.reload(event['args'][0].title()):
                # alert the user that reload was successful
                self.notice(event['nick'],
                    'Successfully reloaded module: %s' % event['args'][0])
                # re-register any commands
                self.module_handler.modules[event['args'][0].title()]._register_events()
            else:
                # alert user that reload failed
                self.notice(event['nick'],
                    'Could not reload module: %s' % event['args'][0])
        else:

            self.syntax_message(event['nick'], '.reload <module>')

    @owner
    def raw(self, event):
        """Processes a raw python command from IRC.

        @param the python command to run

        """
        try:
            print ' '.join(event['args'])
            self.msg(event['target'], eval(' '.join(event['args'])))
        except Exception, e:
            print e

    def welcome(self, event):
        """Event handler for when the bot connects to a server."""

        print "Connected to network: %s." % self.server.server_config.name

    	for channel in self.server.server_config.channels:
    	    self.server.join(channel.name)

<|endoftext|>


<commit_before>
"""
Admin for the grades app
"""
from django.contrib import admin

from exams import models

class ExamRunAdmin(admin.ModelAdmin):
    """Admin for ExamRun"""
    model = models.ExamRun
    list_display = (
        'id',
        'course',
        'semester',
        'exam_series_code',
        'date_first_schedulable',
        'date_last_schedulable',
        'date_first_eligible',
        'date_last_eligible',
        'authorized',
    )
    list_filter = ('course__title', 'course__program__title', 'semester', )
    ordering = ('-date_first_eligible',)
    readonly_fields = ('authorized',)

    def get_readonly_fields(self, request, obj=None):
        """Conditionally determine readonly fields"""
        if not self.is_modifiable(obj):
            # exam_series_code cannot be changed due to Pearson requirement
            return self.readonly_fields + ('exam_series_code',)
        return self.readonly_fields

    def has_delete_permission(self, request, obj=None):
        """Whether record can be deleted or not"""
        return self.is_modifiable(obj)

    def is_modifiable(self, exam_run):
        """
        Determines if an ExamRun can be modified/deleted

        Returns:
            bool: True if the run can be modified/deleted
        """
        return exam_run is None or exam_run.id is None or not exam_run.has_authorizations

admin.site.register(models.ExamRun, ExamRunAdmin)
<commit_msg>adding admin panel entry for ExamAuthorization (#4724)<commit_after>
"""
Admin for the grades app
"""
from django.contrib import admin

from exams import models

class ExamRunAdmin(admin.ModelAdmin):
    """Admin for ExamRun"""
    model = models.ExamRun
    list_display = (
        'id',
        'course',
        'semester',
        'exam_series_code',
        'date_first_schedulable',
        'date_last_schedulable',
        'date_first_eligible',
        'date_last_eligible',
        'authorized',
    )
    list_filter = ('course__title', 'course__program__title', 'semester', )
    ordering = ('-date_first_eligible',)
    readonly_fields = ('authorized',)

    def get_readonly_fields(self, request, obj=None):
        """Conditionally determine readonly fields"""
        if not self.is_modifiable(obj):
            # exam_series_code cannot be changed due to Pearson requirement
            return self.readonly_fields + ('exam_series_code',)
        return self.readonly_fields

    def has_delete_permission(self, request, obj=None):
        """Whether record can be deleted or not"""
        return self.is_modifiable(obj)

    def is_modifiable(self, exam_run):
        """
        Determines if an ExamRun can be modified/deleted

        Returns:
            bool: True if the run can be modified/deleted
        """
        return exam_run is None or exam_run.id is None or not exam_run.has_authorizations

class ExamAuthorizationAdmin(admin.ModelAdmin):
    """Admin for ExamAuthorization"""

    model = models.ExamAuthorization
    list_display = (
        'id',
        'user_email',
        'course_number',
        'exam_run_id',
        'exam_coupon_url',
    )
    list_filter = (
        'exam_run__id',
        'course__course_number',
        'course__title',
    )
    search_fields = (
        'exam_run__id',
        'course__course_number',
    )
    raw_id_fields = ('user',)

    def user_email(self, obj):
        """Getter for the User foreign-key element email"""
        return obj.user.email

    def course_number(self, obj):
        """Getter for the Course foreign-key element course_number"""
        return obj.course.course_number

    def exam_run_id(self, obj):
        """Getter for the ExamRun foreign-key element id"""
        return obj.exam_run.id

admin.site.register(models.ExamRun, ExamRunAdmin)
admin.site.register(models.ExamAuthorization, ExamAuthorizationAdmin)
<|endoftext|>


<commit_before>#!/usr/bin/python
from datetime import datetime
from fcntl import fcntl, F_GETFL, F_SETFL
import logging
from logging import handlers
from os import O_NONBLOCK
import re
import subprocess
import time
from subprocess import PIPE, STDOUT

"""
Standalone daemon meant to be run to poll weather stations in parallel.
Example to run using uwsgi:

    attach-daemon2  = exec=venv/bin/python pollstations.py

Should be run from dir with venv and manage.py available.
"""

# Ensure that the logs dir is present.
subprocess.call(['mkdir', '-p', 'logs'])
# Set up logging in a standardised way.
logger = logging.getLogger('pollstations')
logger.setLevel(logging.INFO)
fh = handlers.RotatingFileHandler(
    'logs/pollstations.log', maxBytes=5 * 1024 * 1024, backupCount=5)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
fh.setFormatter(formatter)
logger.addHandler(fh)

def get_stations():
    """Query the weather app for active AW stations and return a dict of
    station metadata. The dict is used to manage polling processes and the
    interval of polling. Format of the returned dict:

    {
        'IP_ADDRESS': {
            'port': PORT_NUMBER_INT,
            'interval': MINUTES_INT,
            'polled': None,
            'process': None
        },
        ...
    }
    """
    # Parse the list of IPs for active weather stations.
    try:
        station_string = subprocess.check_output(
            ['venv/bin/python', 'manage.py', 'station_metadata'], stderr=STDOUT)
    except subprocess.CalledProcessError as e:
        logger.error(e.output)
        return False
    stations = {}
    # Generate a dict of station IP, port, interval, last polled and process.
    for i in [s for s in station_string.strip().split(",")]:
        if i:
            station = i.split(':')
            stations[station[0]] = {
                'port': station[1],
                'interval': int(station[2]),
                'polled': None,
                'process': None,
                'process_start': None
            }
    return stations

def connect_station(ip, port):
    """Connect to a weather station via Telnet using subprocess, in a
    non-blocking manner.
    """
    logger.info("Connecting to station {}:{}".format(ip, port))
    p = subprocess.Popen(["/usr/bin/telnet", ip, port], stdin=PIPE, stdout=PIPE, stderr=PIPE)
    flags = fcntl(p.stdout, F_GETFL)
    fcntl(p.stdout, F_SETFL, flags | O_NONBLOCK)
    return p

# Get the dict of weather stations.
try:
    STATIONS = get_stations()
    if STATIONS:
        polling = True
    else:
        polling = False  # No active weather stations.
except subprocess.CalledProcessError as e:
    # We can't do anything without this dict, so abort.
    polling = False

# Observation string patterns
OBS_PATTERNS = [
    '([A-Z]+=\d*\.?\d*\|)',  # Telvent
    '(^0R0),([A-Za-z]{2}=\d+[A-Z])',  # Vaisala
]

while polling:
    # For each station, instantiate a process to poll it for observation data
    # at defined intervals (in minutes).
    # We check to see if the required polling interval has passed and if so, we
    # create a process to poll the station immediately for an observation.
    for ip, metadata in STATIONS.iteritems():
        port = metadata['port']
        process = metadata['process']
        started = metadata['process_start']
        interval = metadata['interval']
        last_poll = metadata['polled']
        now = datetime.now()
        # If the polling interval has passed, create a process and read data.
        if not last_poll or interval <= 1 or ((now - last_poll).total_seconds() / 60) >= interval:
            # No existing process/terminated process: create one now.
            if not process or process.poll() is not None:  # Non-existent/terminated.
                # Establish a connection.
                STATIONS[ip]['process'] = process = connect_station(ip, port)
                STATIONS[ip]['process_start'] = datetime.now()
            try:
                # Read up to 2kb of the connection output and parse the
                # observation string from that.
                data = process.stdout.read(2048)
                # Split the lines of the response and find the observation.
                for line in data.strip().split('\n'):
                    line = line.strip()
                    for pattern in OBS_PATTERNS:
                        if re.search(pattern, line):  # Found matching pattern.
                            STATIONS[ip]['polled'] = datetime.now()
                            # Looks like an observation string; save it.
                            obs = '{}::{}'.format(ip, line)
                            logger.info('Observation data: {}'.format(obs))
                            # This mgmt command will write the observation to the
                            # database and optionally upload it to DAFWA.
                            try:
                                output = subprocess.check_output([
                                    'venv/bin/python', 'manage.py', 'write_observation', obs])
                                logger.info(output.strip())
                            except subprocess.CalledProcessError as e:
                                logger.error(e.output)
                            # Terminate this polling process, it's finished with.
                            age = (now - started).total_seconds()
                            logger.info('IP {} poller PID {} killed at {} seconds old'.format(ip, process.pid, age))
                            process.kill()
            except Exception as e:
                continue
    # Pause, and repeat the polling loop.
    time.sleep(3)
<commit_msg>pollstations: terminate >1 min hung process<commit_after>#!/usr/bin/python
from datetime import datetime
from fcntl import fcntl, F_GETFL, F_SETFL
import logging
from logging import handlers
from os import O_NONBLOCK
import re
import subprocess
import time
from subprocess import PIPE, STDOUT

"""
Standalone daemon meant to be run to poll weather stations in parallel.
Example to run using uwsgi:

    attach-daemon2  = exec=venv/bin/python pollstations.py

Should be run from dir with venv and manage.py available.
"""

# Ensure that the logs dir is present.
subprocess.call(['mkdir', '-p', 'logs'])
# Set up logging in a standardised way.
logger = logging.getLogger('pollstations')
logger.setLevel(logging.INFO)
fh = handlers.RotatingFileHandler(
    'logs/pollstations.log', maxBytes=5 * 1024 * 1024, backupCount=5)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
fh.setFormatter(formatter)
logger.addHandler(fh)

def get_stations():
    """Query the weather app for active AW stations and return a dict of
    station metadata. The dict is used to manage polling processes and the
    interval of polling. Format of the returned dict:

    {
        'IP_ADDRESS': {
            'port': PORT_NUMBER_INT,
            'interval': MINUTES_INT,
            'polled': None,
            'process': None
        },
        ...
    }
    """
    # Parse the list of IPs for active weather stations.
    try:
        station_string = subprocess.check_output(
            ['venv/bin/python', 'manage.py', 'station_metadata'], stderr=STDOUT)
    except subprocess.CalledProcessError as e:
        logger.error(e.output)
        return False
    stations = {}
    # Generate a dict of station IP, port, interval, last polled and process.
    for i in [s for s in station_string.strip().split(",")]:
        if i:
            station = i.split(':')
            stations[station[0]] = {
                'port': station[1],
                'interval': int(station[2]),
                'polled': None,
                'process': None,
                'process_start': None
            }
    return stations

def connect_station(ip, port):
    """Connect to a weather station via Telnet using subprocess, in a
    non-blocking manner.
    """
    logger.info("Connecting to station {}:{}".format(ip, port))
    p = subprocess.Popen(["/usr/bin/telnet", ip, port], stdin=PIPE, stdout=PIPE, stderr=PIPE)
    flags = fcntl(p.stdout, F_GETFL)
    fcntl(p.stdout, F_SETFL, flags | O_NONBLOCK)
    return p

# Get the dict of weather stations.
try:
    STATIONS = get_stations()
    if STATIONS:
        polling = True
    else:
        polling = False  # No active weather stations.
except subprocess.CalledProcessError as e:
    # We can't do anything without this dict, so abort.
    polling = False

# Observation string patterns
OBS_PATTERNS = [
    '([A-Z]+=\d*\.?\d*\|)',  # Telvent
    '(^0R0),([A-Za-z]{2}=\d+[A-Z])',  # Vaisala
]

while polling:
    # For each station, instantiate a process to poll it for observation data
    # at defined intervals (in minutes).
    # We check to see if the required polling interval has passed and if so, we
    # create a process to poll the station immediately for an observation.
    for ip, metadata in STATIONS.iteritems():
        port = metadata['port']
        process = metadata['process']
        started = metadata['process_start']
        interval = metadata['interval']
        last_poll = metadata['polled']
        now = datetime.now()
        # If the polling interval has passed, create a process and read data.
        if not last_poll or interval <= 1 or ((now - last_poll).total_seconds() / 60) >= interval:
            # No existing process/terminated process: create one now.
            if not process or process.poll() is not None:  # Non-existent/terminated.
                # Establish a connection.
                STATIONS[ip]['process'] = process = connect_station(ip, port)
                STATIONS[ip]['process_start'] = datetime.now()
            try:
                # Read up to 2kb of the connection output and parse the
                # observation string from that.
                data = process.stdout.read(2048)
                # Split the lines of the response and find the observation.
                for line in data.strip().split('\n'):
                    line = line.strip()
                    for pattern in OBS_PATTERNS:
                        if re.search(pattern, line):  # Found matching pattern.
                            STATIONS[ip]['polled'] = datetime.now()
                            # Looks like an observation string; save it.
                            obs = '{}::{}'.format(ip, line)
                            logger.info('Observation data: {}'.format(obs))
                            # This mgmt command will write the observation to the
                            # database and optionally upload it to DAFWA.
                            try:
                                output = subprocess.check_output([
                                    'venv/bin/python', 'manage.py', 'write_observation', obs])
                                logger.info(output.strip())
                            except subprocess.CalledProcessError as e:
                                logger.error(e.output)
                            # Terminate this polling process if >1 minute, it's finished with.
                            if interval > 1:
                                age = (now - started).total_seconds()
                                logger.info('IP {} poller PID {} killed at {} seconds old'.format(ip, process.pid, age))
                                process.kill()
            except Exception as e:
                continue
    # Pause, and repeat the polling loop.
    time.sleep(3)
<|endoftext|>


<commit_before>#! /usr/bin/python
import argparse
from datetime import datetime
import sys

import PIL.Image

from .polygrid import PolyGrid
from .shapes import supershapes_dict
from .maze import Maze

ss_dict = supershapes_dict()

def commandline():
    parser = _parser()
    kwargs = vars(parser.parse_args())
    # setup the base grid with a supershape if provided
    grid = PolyGrid(supershape=ss_dict.get(kwargs.pop('shape'), None))

    # pull off non-common parameters
    text = kwargs.pop('text')
    if text is not None:
        text = text.decode(sys.stdin.encoding)
    image_path = kwargs.pop('image')
    if image_path is not None:
        image_path = image_path.decode(sys.stdin.encoding)
    font_path = kwargs.pop('font')
    if font_path is not None:
        font_path = font_path.decode(sys.stdin.encoding)
    filename = kwargs.pop('output')
    if filename is not None:
        filename = filename.decode(sys.stdin.encoding)

    # fill the grid and create maze based on the remaining arguments provided
    if text:
        grid.create_string(text, font_path=font_path, **kwargs)
        maze_type = 'Text'
    elif image_path:
        image = PIL.Image.open(image_path).convert('L')
        if not image:
            print('Unable to open the provided image path: {}'
                  ''.format(image_path))
        grid.create_from_image(image, **kwargs)
        maze_type = 'Image'
    else:
        grid.create_rectangle(**kwargs)
        maze_type = 'Rectangle'

    maze = Maze(grid)
    save_maze(maze, maze_type, filename)

def save_maze(maze, maze_type, filename=None):
    image = maze.image()
    if image is None:
        print('This maze appears to be empty. Not saving.')
    else:
        now_str = str(datetime.now().time())
        clean_now_string = now_str.replace(':', '.').rsplit('.', 1)[0]
        if filename is None:
            filename = '{} - {} made with {}'.format(clean_now_string,
                                                     maze_type,
                                                     maze.shape_name())
        # force png... for your own good! png works well with this type
        # of image, lossless AND smaller file size than jpg.
        # If this is ever updated to work with original images remaining
        # in the background of the maze, then jpg might make sense.
        filename += '.png'
        image.save(filename)
        print(u'Saved {}'.format(filename))

def _parser():
    parser = argparse.ArgumentParser(description='Make and save mazes.')
    # optional top level type of maze to make (default rectangle)
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-t', '--text',
                       help='Make a maze inside the characters of TEXT.'
                            ' Note: interprets "\\n" as a new line.')
    group.add_argument('-i', '--image',
                       help='Make a maze from IMAGE (path).')
    # optional complexity
    parser.add_argument('-c', '--complexity', type=_positive,
                        help='Positive number that controls maze complexity.'
                             ' For example: 0.5 is very easy. 110 is very hard.')
     # optional shape to use
    ss_names = ss_dict.keys()
    parser.add_argument('-s', '--shape', choices=ss_names,
                        help='Make the maze with this shape. Random otherwise.')
    # optional aspect (relative height)
    parser.add_argument('-a', '--aspect', type=_positive,
                        help='Set the height/width aspect of the maze.')
    parser.add_argument('-f', '--font', type=str,
                        help='Provide a font path for text mazes.')
    parser.add_argument('-o', '--output', type=str,
                        help='Output filename. The format will always be PNG.')
    return parser

def _positive(v):
    try:
        number = float(v)
    except Exception:
        raise argparse.ArgumentTypeError('{} is not a valid number'.format(v))
    if number <= 0:
        raise argparse.ArgumentTypeError('{} must be positive'.format(number))
    return number

if __name__ == '__main__':
    commandline()
<commit_msg>Update cli to work with new size options<commit_after>#! /usr/bin/python
import argparse
from datetime import datetime
import sys

import PIL.Image

from .polygrid import PolyGrid
from .shapes import supershapes_dict
from .maze import Maze

ss_dict = supershapes_dict()

def commandline():
    parser = _parser()
    kwargs = vars(parser.parse_args())
    # setup the base grid with a supershape if provided
    grid = PolyGrid(supershape=ss_dict.get(kwargs.pop('shape'), None))

    # pull off non-common parameters
    text = kwargs.pop('text')
    if text is not None:
        text = text.decode(sys.stdin.encoding)
    image_path = kwargs.pop('image')
    if image_path is not None:
        image_path = image_path.decode(sys.stdin.encoding)
    font_path = kwargs.pop('font')
    if font_path is not None:
        font_path = font_path.decode(sys.stdin.encoding)
    filename = kwargs.pop('output')
    if filename is not None:
        filename = filename.decode(sys.stdin.encoding)

    # fill the grid and create maze based on the remaining arguments provided
    if text:
        grid.create_string(text, font_path=font_path, **kwargs)
        maze_type = 'Text'
    elif image_path:
        image = PIL.Image.open(image_path).convert('L')
        if not image:
            print('Unable to open the provided image path: {}'
                  ''.format(image_path))
        grid.create_from_image(image, **kwargs)
        maze_type = 'Image'
    else:
        grid.create_rectangle(**kwargs)
        maze_type = 'Rectangle'

    maze = Maze(grid)
    save_maze(maze, maze_type, filename)

def save_maze(maze, maze_type, filename=None):
    image = maze.image()
    if image is None:
        print('This maze appears to be empty. Not saving.')
    else:
        now_str = str(datetime.now().time())
        clean_now_string = now_str.replace(':', '.').rsplit('.', 1)[0]
        if filename is None:
            filename = '{} - {} made with {}'.format(clean_now_string,
                                                     maze_type,
                                                     maze.shape_name())
        # force png... for your own good! png works well with this type
        # of image, lossless AND smaller file size than jpg.
        # If this is ever updated to work with original images remaining
        # in the background of the maze, then jpg might make sense.
        filename += '.png'
        image.save(filename)
        print(u'Saved {}'.format(filename))

def _parser():
    parser = argparse.ArgumentParser(description='Make and save mazes.')
    # optional top level type of maze to make
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-t', '--text',
                       help='Make a maze inside the characters of TEXT.'
                            ' Note: interprets "\\n" as a new line.')
    group.add_argument('-i', '--image',
                       help='Make a maze from IMAGE (path).')

    # optional complexity, size and aspect arguments.
    group = parser.add_argument_group(title='Optional complexity, size and aspect',
                                      description='Note: If over-specified, the choices are resolved internally.')
    group.add_argument('-c', '--complexity', type=_positive,
                       help='Positive number that controls maze complexity.'
                            ' For example: 0.5 is very easy. 110 is very hard.')
    group.add_argument('-w', '--width', type=_positive,
                       help='Positive number that sets the width of the maze in side-length units.')
    group.add_argument('-g', '--height', type=_positive,
                       help='Positive number that sets the height of the maze in side-length units.')
    group.add_argument('-a', '--aspect', type=_positive,
                       help='Set the height/width aspect of the maze.')

    # optional shape to use
    ss_names = ss_dict.keys()
    parser.add_argument('-s', '--shape', choices=ss_names,
                        help='Make the maze with this shape. Random otherwise.')
    # optional aspect (relative height)
    parser.add_argument('-f', '--font', type=str,
                        help='Provide a font path for text mazes.')
    parser.add_argument('-o', '--output', type=str,
                        help='Output filename. The format will always be PNG.')
    return parser

def _positive(v):
    try:
        number = float(v)
    except Exception:
        raise argparse.ArgumentTypeError('{} is not a valid number'.format(v))
    if number <= 0:
        raise argparse.ArgumentTypeError('{} must be positive'.format(number))
    return number

if __name__ == '__main__':
    commandline()
<|endoftext|>


<commit_before>"""Tests for acme.dns_resolver."""
import unittest

import dns
import mock

from acme import dns_resolver

class TxtRecordsForNameTest(unittest.TestCase):

    def create_txt_response(self, name, txt_records):
        """
        Returns an RRSet containing the 'txt_records' as the result of a DNS
        query for 'name'.

        This takes advantage of the fact that an Answer object mostly behaves
        like an RRset.
        """
        return dns.rrset.from_text_list(name, 60, "IN", "TXT", txt_records)

    @mock.patch("acme.dns_resolver.dns.resolver.query")
    def test_txt_records_for_name_test_with_single_response(self, mock_dns):
        mock_dns.return_value = self.create_txt_response('name', ['response'])
        self.assertEqual(['response'],
                         dns_resolver.txt_records_for_name('name'))

    @mock.patch("acme.dns_resolver.dns.resolver.query")
    def test_txt_records_for_name_with_multiple_responses(self, mock_dns):
        mock_dns.return_value = self.create_txt_response(
            'name', ['response1', 'response2'])
        self.assertEqual(['response1', 'response2'],
                         dns_resolver.txt_records_for_name('name'))

    @mock.patch("acme.dns_resolver.dns.resolver.query")
    def test_txt_records_for_name_domain_not_found(self, mock_dns):
        mock_dns.side_effect = dns.exception.DNSException
        self.assertEquals([], dns_resolver.txt_records_for_name('name'))
<commit_msg>Skip dns_resolver tests if dnspython is not available.<commit_after>"""Tests for acme.dns_resolver."""
import unittest

import mock

try:
    import dns
except ImportError:
    dns = None

from acme import dns_resolver

@unittest.skipIf(dns is None,
                 "dnspython is not available, skipping dns_resolver tests")
class TxtRecordsForNameTest(unittest.TestCase):

    def create_txt_response(self, name, txt_records):
        """
        Returns an RRSet containing the 'txt_records' as the result of a DNS
        query for 'name'.

        This takes advantage of the fact that an Answer object mostly behaves
        like an RRset.
        """
        return dns.rrset.from_text_list(name, 60, "IN", "TXT", txt_records)

    @mock.patch("acme.dns_resolver.dns.resolver.query")
    def test_txt_records_for_name_test_with_single_response(self, mock_dns):
        mock_dns.return_value = self.create_txt_response('name', ['response'])
        self.assertEqual(['response'],
                         dns_resolver.txt_records_for_name('name'))

    @mock.patch("acme.dns_resolver.dns.resolver.query")
    def test_txt_records_for_name_with_multiple_responses(self, mock_dns):
        mock_dns.return_value = self.create_txt_response(
            'name', ['response1', 'response2'])
        self.assertEqual(['response1', 'response2'],
                         dns_resolver.txt_records_for_name('name'))

    @mock.patch("acme.dns_resolver.dns.resolver.query")
    def test_txt_records_for_name_domain_not_found(self, mock_dns):
        mock_dns.side_effect = dns.exception.DNSException
        self.assertEquals([], dns_resolver.txt_records_for_name('name'))
<|endoftext|>


<commit_before>from nose.tools import assert_equal
import epics
p = None

def test_epics_smoke():

    print(p)
    pv = epics.PV('BSTEST:VAL')
    pv.value = 123456
    print(pv)
    print(pv.value)
    print(pv.connect())
    assert pv.connect()
    for j in range(1, 15):
        pv.put(j, wait=True)
        ret = pv.get(use_monitor=False)
        assert_equal(ret, j)
<commit_msg>TST: test coverage from epics_callbacks<commit_after>from functools import partial
from nose.tools import assert_equal
import asyncio
import epics
import time as ttime

from bluesky import Msg
from bluesky.epics_callbacks import (PVSuspendBoolHigh,
                                     PVSuspendBoolLow,
                                     PVSuspendFloor,
                                     PVSuspendCeil,
                                     PVSuspendInBand,
                                     PVSuspendOutBand)

from bluesky.tests.utils import setup_test_run_engine
RE = setup_test_run_engine()
loop = asyncio.get_event_loop()

def test_epics_smoke():

    pv = epics.PV('BSTEST:VAL')
    pv.value = 123456
    print(pv)
    print(pv.value)
    print(pv.connect())
    assert pv.connect()
    for j in range(1, 15):
        pv.put(j, wait=True)
        ret = pv.get(use_monitor=False)
        assert_equal(ret, j)

def _test_suspender(suspender_class, sc_args, start_val, fail_val,
                    resume_val, wait_time):

    my_suspender = suspender_class(RE, 'BSTEST:VAL', *sc_args, sleep=wait_time)
    print(my_suspender._lock)
    pv = epics.PV('BSTEST:VAL')
    putter = partial(pv.put, wait=True)
    # make sure we start at good value!
    putter(start_val)
    # dumb scan
    scan = [Msg('checkpoint'), Msg('sleep', None, .2)]
    # paranoid
    assert_equal(RE.state, 'idle')

    start = ttime.time()
    # queue up fail and resume conditions
    loop.call_later(.1, putter, fail_val)
    loop.call_later(1, putter, resume_val)
    # start the scan
    RE(scan)
    stop = ttime.time()
    # paranoid clean up of pv call back
    my_suspender._pv.disconnect()
    # assert we waited at least 2 seconds + the settle time
    print(stop - start)
    assert stop - start > 1 + wait_time + .2

def test_suspending():
    yield _test_suspender, PVSuspendBoolHigh, (), 0, 1, 0, .5
    yield _test_suspender, PVSuspendBoolLow, (), 1, 0, 1, .5
    yield _test_suspender, PVSuspendFloor, (.5,), 1, 0, 1, .5
    yield _test_suspender, PVSuspendCeil, (.5,), 0, 1, 0, .5
    yield _test_suspender, PVSuspendInBand, (.5, 1.5), 1, 0, 1, .5
    yield _test_suspender, PVSuspendOutBand, (.5, 1.5), 0, 1, 0, .5
<|endoftext|>


<commit_before>import os
import sys
import traceback
from multiprocessing import Process, Manager, Queue, cpu_count, Value, Lock, JoinableQueue
import time
from queue import Empty, Full
from collections import OrderedDict

from numpy import zeros

from functools import partial

from acousticsim.representations import Envelopes, Mfcc

from acousticsim.distance import dtw_distance, xcorr_distance, dct_distance

from acousticsim.exceptions import AcousticSimError,NoWavError, AcousticSimPythonError

def default_njobs():
    return int(0.75 * cpu_count())

class Counter(object):
    def __init__(self, initval=0):
        self.val = Value('i', initval)
        self.lock = Lock()

    def increment(self):
        with self.lock:
            self.val.value += 1

    def value(self):
        with self.lock:
            return self.val.value

class Stopped(object):
    def __init__(self, initval=False):
        self.val = Value('i', initval)
        self.lock = Lock()

    def stop(self):
        with self.lock:
            self.val.value = True

    def stop_check(self):
        with self.lock:
            return self.val.value

class RepWorker(Process):
    def __init__(self, job_q, return_dict, rep_func, attributes, counter, stopped):
        Process.__init__(self)
        self.job_q = job_q
        self.return_dict = return_dict
        self.function = rep_func
        self.attributes = attributes
        self.counter = counter
        self.stopped = stopped

    def run(self):
        while True:
            self.counter.increment()
            try:
                filename = self.job_q.get(timeout=1)
            except Empty:
                break
            self.job_q.task_done()
            if self.stopped.stop_check():
                continue
            if not os.path.exists(filename):
                continue
            path, filelabel = os.path.split(filename)

            att = OrderedDict()
            att['filename'] = filelabel
            try:
                att.update(self.attributes[filelabel])
            except (KeyError, TypeError):
                pass
            true_label = os.path.split(path)[1]
            try:
                rep = self.function(filename,attributes=att)
                rep._true_label = true_label
                self.return_dict[os.path.split(filename)[1]] = rep
            except Exception as e:
                self.stopped.stop()
                self.return_dict['error'] = AcousticSimPythonError(traceback.format_exception(*sys.exc_info()))

        return

class DistWorker(Process):
    def __init__(self, job_q, return_dict, counter, dist_func, output_sim, axb, cache, stopped):
        Process.__init__(self)
        self.job_q = job_q
        self.return_dict = return_dict
        self.function = dist_func
        self.output_sim = output_sim
        self.axb = axb
        self.cache = cache
        self.counter = counter
        self.stopped = stopped

    def run(self):
        while True:
            self.counter.increment()
            try:
                pm = self.job_q.get(timeout=1)
            except Empty:
                break
            self.job_q.task_done()
            if self.stopped.stop_check():
                continue
            filetup = tuple(map(lambda x: os.path.split(x)[1],pm))
            try:
                base = self.cache[filetup[0]]
                model = self.cache[filetup[1]]
                if self.axb:
                    shadow = self.cache[filetup[2]]
            except KeyError:
                continue
            try:
                dist1 = self.function(base,model)
                if self.axb:
                    dist2 = self.function(shadow,model)
                    try:
                        ratio = dist2 / dist1
                    except ZeroDivisionError:
                        ratio = -1
                else:
                    ratio = dist1
                if self.output_sim:
                    try:
                        ratio = 1/ratio
                    except ZeroDivisionError:
                        ratio = 1
                self.return_dict[filetup] = ratio
            except Exception as e:
                self.stopped.stop()
                self.return_dict['error'] = AcousticSimPythonError(traceback.format_exception(*sys.exc_info()))

        return

def generate_cache(path_mapping,rep_func, attributes,num_procs, call_back, stop_check):

    all_files = set()
    for pm in path_mapping:
        if stop_check is not None and stop_check():
            return
        all_files.update(pm)
    all_files = sorted(all_files)
    stopped = Stopped()
    job_queue = JoinableQueue(100)
    file_ind = 0
    while True:
        if file_ind == len(all_files):
            break
        try:
            job_queue.put(all_files[file_ind],False)
        except Full:
            break
        file_ind += 1
    manager = Manager()
    return_dict = manager.dict()
    procs = []

    counter = Counter()
    for i in range(num_procs):
        p = RepWorker(job_queue,
                      return_dict,rep_func,attributes, counter, stopped)
        procs.append(p)
        p.start()
    if call_back is not None:
        call_back('Generating representations...')

    while True:
        if file_ind == len(all_files):
            break
        if stop_check is not None and stop_check():
            stopped.stop()
            time.sleep(1)
            break
        job_queue.put(all_files[file_ind])

        if call_back is not None:
            value = counter.value()
            call_back(value)
        file_ind += 1
    job_queue.join()

    for p in procs:
        p.join()
    if 'error' in return_dict:
        raise(return_dict['error'])
    return return_dict

def calc_asim(path_mapping, cache,dist_func, output_sim,num_procs, call_back, stop_check):
    if len(path_mapping[0]) == 3:
        axb = True
    else:
        axb = False
    job_queue = JoinableQueue(100)
    stopped = Stopped()

    map_ind = 0
    while True:
        if map_ind == len(path_mapping):
            break
        try:
            job_queue.put(path_mapping[map_ind],False)
        except Full:
            break
        map_ind += 1

    manager = Manager()
    return_dict = manager.dict()
    procs = []

    counter = Counter()
    for i in range(num_procs):
        p = DistWorker(job_queue,
                      return_dict, counter, dist_func, output_sim, axb, cache, stopped)
        procs.append(p)
        p.start()

    if call_back is not None:
        call_back('Calculating acoustic similarity...')
        call_back(0,len(path_mapping))

    while True:
        if map_ind == len(path_mapping):
            break
        if stop_check is not None and stop_check():
            stopped.stop()
            time.sleep(1)
            break
        job_queue.put(path_mapping[map_ind])

        if call_back is not None:
            value = counter.value()
            call_back(value)
        map_ind += 1

    job_queue.join()

    for p in procs:
        p.join()
    if 'error' in return_dict:
        raise(return_dict['error'])
    return return_dict

<commit_msg>Add support for generator path mappings<commit_after>import os
import sys
import traceback
from multiprocessing import Process, Manager, Queue, cpu_count, Value, Lock, JoinableQueue
import time
from queue import Empty, Full
from collections import OrderedDict

from numpy import zeros

from functools import partial

from acousticsim.representations import Envelopes, Mfcc

from acousticsim.distance import dtw_distance, xcorr_distance, dct_distance

from acousticsim.exceptions import AcousticSimError,NoWavError, AcousticSimPythonError

def default_njobs():
    return int(0.75 * cpu_count())

class Counter(object):
    def __init__(self, initval=0):
        self.val = Value('i', initval)
        self.lock = Lock()

    def increment(self):
        with self.lock:
            self.val.value += 1

    def value(self):
        with self.lock:
            return self.val.value

class Stopped(object):
    def __init__(self, initval=False):
        self.val = Value('i', initval)
        self.lock = Lock()

    def stop(self):
        with self.lock:
            self.val.value = True

    def stop_check(self):
        with self.lock:
            return self.val.value

class RepWorker(Process):
    def __init__(self, job_q, return_dict, rep_func, attributes, counter, stopped):
        Process.__init__(self)
        self.job_q = job_q
        self.return_dict = return_dict
        self.function = rep_func
        self.attributes = attributes
        self.counter = counter
        self.stopped = stopped

    def run(self):
        while True:
            self.counter.increment()
            try:
                filename = self.job_q.get(timeout=1)
            except Empty:
                break
            self.job_q.task_done()
            if self.stopped.stop_check():
                continue
            if not os.path.exists(filename):
                continue
            path, filelabel = os.path.split(filename)

            att = OrderedDict()
            att['filename'] = filelabel
            try:
                att.update(self.attributes[filelabel])
            except (KeyError, TypeError):
                pass
            true_label = os.path.split(path)[1]
            try:
                rep = self.function(filename,attributes=att)
                rep._true_label = true_label
                self.return_dict[os.path.split(filename)[1]] = rep
            except Exception as e:
                self.stopped.stop()
                self.return_dict['error'] = AcousticSimPythonError(traceback.format_exception(*sys.exc_info()))

        return

class DistWorker(Process):
    def __init__(self, job_q, return_dict, counter, dist_func, output_sim, axb, cache, stopped):
        Process.__init__(self)
        self.job_q = job_q
        self.return_dict = return_dict
        self.function = dist_func
        self.output_sim = output_sim
        self.axb = axb
        self.cache = cache
        self.counter = counter
        self.stopped = stopped

    def run(self):
        while True:
            self.counter.increment()
            try:
                pm = self.job_q.get(timeout=1)
            except Empty:
                break
            self.job_q.task_done()
            if self.stopped.stop_check():
                continue
            filetup = tuple(map(lambda x: os.path.split(x)[1],pm))
            try:
                base = self.cache[filetup[0]]
                model = self.cache[filetup[1]]
                if self.axb:
                    shadow = self.cache[filetup[2]]
            except KeyError:
                continue
            try:
                dist1 = self.function(base,model)
                if self.axb:
                    dist2 = self.function(shadow,model)
                    try:
                        ratio = dist2 / dist1
                    except ZeroDivisionError:
                        ratio = -1
                else:
                    ratio = dist1
                if self.output_sim:
                    try:
                        ratio = 1/ratio
                    except ZeroDivisionError:
                        ratio = 1
                self.return_dict[filetup] = ratio
            except Exception as e:
                self.stopped.stop()
                self.return_dict['error'] = AcousticSimPythonError(traceback.format_exception(*sys.exc_info()))

        return

def generate_cache(path_mapping,rep_func, attributes,num_procs, call_back, stop_check):

    all_files = set()
    for pm in path_mapping:
        if stop_check is not None and stop_check():
            return
        all_files.update(pm)
    all_files = sorted(all_files)
    stopped = Stopped()
    job_queue = JoinableQueue(100)
    file_ind = 0
    while True:
        if file_ind == len(all_files):
            break
        try:
            job_queue.put(all_files[file_ind],False)
        except Full:
            break
        file_ind += 1
    manager = Manager()
    return_dict = manager.dict()
    procs = []

    counter = Counter()
    for i in range(num_procs):
        p = RepWorker(job_queue,
                      return_dict,rep_func,attributes, counter, stopped)
        procs.append(p)
        p.start()
    if call_back is not None:
        call_back('Generating representations...')

    while True:
        if file_ind == len(all_files):
            break
        if stop_check is not None and stop_check():
            stopped.stop()
            time.sleep(1)
            break
        job_queue.put(all_files[file_ind])

        if call_back is not None:
            value = counter.value()
            call_back(value)
        file_ind += 1
    job_queue.join()

    for p in procs:
        p.join()
    if 'error' in return_dict:
        raise(return_dict['error'])
    return return_dict

def calc_asim(path_mapping, cache,dist_func, output_sim,num_procs, call_back, stop_check):
    try:
        if len(path_mapping[0]) == 3:
            axb = True
        else:
            axb = False
        mapping_gen = False
    except TypeError:
        axb = False
        mapping_gen = True
    job_queue = JoinableQueue(100)
    stopped = Stopped()
    map_ind = 0
    if mapping_gen:
        while True:
            try:
                pathmap = next(path_mapping)
                map_ind += 1
            except StopIteration:
                break
            try:
                job_queue.put(pathmap,False)
            except Full:
                break

    else:
        while True:
            if map_ind == len(path_mapping):
                break
            try:
                job_queue.put(path_mapping[map_ind],False)
            except Full:
                break
            map_ind += 1

    manager = Manager()
    return_dict = manager.dict()
    procs = []

    counter = Counter()
    for i in range(num_procs):
        p = DistWorker(job_queue,
                      return_dict, counter, dist_func, output_sim, axb, cache, stopped)
        procs.append(p)
        p.start()

    if call_back is not None:
        call_back('Calculating acoustic similarity...')
        try:
            call_back(0,len(path_mapping))
        except TypeError:
            pass

    if mapping_gen:
        while True:
            try:
                pathmap = next(path_mapping)
            except StopIteration:
                break
            if stop_check is not None and stop_check():
                stopped.stop()
                time.sleep(1)
                break
            job_queue.put(pathmap)
            if call_back is not None:
                value = counter.value()
                call_back(value)
            map_ind += 1

    else:
        while True:
            if map_ind == len(path_mapping):
                break
            if stop_check is not None and stop_check():
                stopped.stop()
                time.sleep(1)
                break
            job_queue.put(path_mapping[map_ind])

            if call_back is not None:
                value = counter.value()
                call_back(value)
            map_ind += 1

    job_queue.join()

    for p in procs:
        p.join()
    if 'error' in return_dict:
        raise(return_dict['error'])
    return return_dict

<|endoftext|>


<commit_before>"""
    Example usage of Memcache class:

    >>> from pony import caching
    >>> cache = caching.Memcache()
    >>> cache.set('k1', 'v1')
    True
    >>> cache.set('k2', 'v2', 10)
    True
    >>> cache.set('k3', 'v3')
    True
    >>> cache.set('k4', 'v4')
    True
    >>> cache.items()
    [('k4', 'v4'), ('k3', 'v3'), ('k2', 'v2'), ('k1', 'v1')]
    >>> cache.get('k2')
    'v2'
    >>> from time import sleep; sleep(10)
    >>> print cache.get('k2')
    None
    >>> cache.items()
    [('k4', 'v4'), ('k3', 'v3'), ('k1', 'v1')]
    >>> cache.add('k5', 'v5')
    True
    >>> cache.add('k5', 'v6')
    False
    >>> cache.items()
    [('k5', 'v5'), ('k4', 'v4'), ('k3', 'v3'), ('k1', 'v1')]
    >>> print cache.incr('x')
    None
    >>> cache.set('x', '100')
    True
    >>> cache.incr('x')
    101
    >>> cache.decr('x', 10)
    91
    >>> cache.get('k3')
    'v3'
    >>> cache.items()
    [('k3', 'v3'), ('x', '91'), ('k5', 'v5'), ('k4', 'v4'), ('k1', 'v1')]
    >>> len(cache)
    5
    >>> 'x' in cache
    True
    >>> cache.replace('k4', 'v4_new')
    True
    >>> cache.replace('notexists', 'value')
    False
    >>> cache.items()
    [('k4', 'v4_new'), ('k3', 'v3'), ('x', '91'), ('k5', 'v5'), ('k1', 'v1')]
"""

from threading import Lock
from time import time

class Node(object):
    __slots__ = 'prev', 'next', 'expire', 'key', 'value'

def normalize(key, value="", expire=None):
    if isinstance(key, tuple): hash_value, key = key
    elif not isinstance(key, str): raise ValueError('Key must be tuple or string. Got: %s' % key.__class__.__name__)
    if not isinstance(value, str): raise ValueError('Value must be string. Got: %s' % value.__class__.__name__)
    if len(value) > 1024 * 1024: raise ValueError('Value size too big: %d' % len(value))
    if expire is not None: expire = time() + expire
    return key, value, expire

class Memcache(object):
    def __init__(self):
        self.dict = {}
        list = self.list = Node()
        list.prev = list.next = list.expire = list.key = list.value = None
        list.prev = list.next = list
        self.lock = Lock()
    def __len__(self):
        return len(self.dict)
    def __contains__(self, key):
        return key in self.dict
    def __iter__(self):
        return iter(self.items())
    def items(self):
        now = time()
        result = []
        append = result.append
        list = self.list
        self.lock.acquire()
        try:
            node = list.prev
            while node is not list:
                expire = node.expire
                if expire is not None and expire <= now:
                    prev, next = node.prev, node.next
                    prev.next = next
                    next.prev = prev
                    del self.dict[node.key]
                else: append((node.key, node.value))
                node = node.prev
        finally: self.lock.release()
        return result
    def _find(self, key, do_create_when_not_present):
        node = self.dict.get(key)
        if node is not None:
            expire = node.expire
            success = expire is None or expire > time()
            prev, next, expire = node.prev, node.next, node.expire
            prev.next = next
            next.prev = prev
        else:
            success = False
            if do_create_when_not_present:
                self.dict[key] = node = Node()
                node.key = key
        if node is not None:
            list = self.list
            node.next = list
            node.prev = prev_top = list.prev
            list.prev = prev_top.next = node
        return success, node
    def get(self, key):
        key, _, _ = normalize(key)
        self.lock.acquire()
        try:
            success, node = self._find(key, False)
            if success: return node.value
            return None
        finally: self.lock.release()
    def set(self, key, value, expire=None):
        key, value, expire = normalize(key, value, expire)
        self.lock.acquire()
        try:
            success, node = self._find(key, True)
            node.value, node.expire = value, expire
        finally: self.lock.release()
        return True
    def add(self, key, value, expire=None):
        key, value, expire = normalize(key, value, expire)
        self.lock.acquire()
        try:
            success, node = self._find(key, True)
            if success: return False
            node.value, node.expire = value, expire
        finally: self.lock.release()
        return True
    def replace(self, key, value, expire=None):
        key, value, expire = normalize(key, value, expire)
        self.lock.acquire()
        try:
            success, node = self._find(key, False)
            if not success: return False
            node.value, node.expire = value, expire
        finally: self.lock.release()
        return True
    def incr(self, key, delta=1):
        key, _, _ = normalize(key)
        self.lock.acquire()
        try:
            success, node = self._find(key, False)
            if not success: return None
            try: value = int(node.value) + delta
            except ValueError: return None
            node.value = str(value)
        finally: self.lock.release()
        return value
    def decr(self, key, delta=1):
        key, _, _ = normalize(key)
        self.lock.acquire()
        try:
            success, node = self._find(key, False)
            if not success: return None
            try: value = int(node.value) - delta
            except ValueError: return None
            node.value = str(value)
        finally: self.lock.release()
        return value
    def flush_all(self):
        self.lock.acquire()
        try:
            self.dict.clear()
            self.list.prev = self.list.next = self.list
            for node in self.dict.itervalues():
                node.prev = node.next = None
        finally: self.lock.release()
<commit_msg>1. Memcache.delete(key) method was added 2. Memcache data cannot exceed max data size specified in Memcache.__init__ call<commit_after>"""
    Example usage of Memcache class:

    >>> from pony import caching
    >>> cache = caching.Memcache()
    >>> cache.set('k1', 'v1')
    True
    >>> cache.set('k2', 'v2', 10)
    True
    >>> cache.set('k3', 'v3')
    True
    >>> cache.set('k4', 'v4')
    True
    >>> cache.items()
    [('k4', 'v4'), ('k3', 'v3'), ('k2', 'v2'), ('k1', 'v1')]
    >>> cache.get('k2')
    'v2'
    >>> from time import sleep; sleep(10)
    >>> print cache.get('k2')
    None
    >>> cache.items()
    [('k4', 'v4'), ('k3', 'v3'), ('k1', 'v1')]
    >>> cache.add('k5', 'v5')
    True
    >>> cache.add('k5', 'v6')
    False
    >>> cache.items()
    [('k5', 'v5'), ('k4', 'v4'), ('k3', 'v3'), ('k1', 'v1')]
    >>> print cache.incr('x')
    None
    >>> cache.set('x', '100')
    True
    >>> cache.incr('x')
    101
    >>> cache.decr('x', 10)
    91
    >>> cache.get('k3')
    'v3'
    >>> cache.items()
    [('k3', 'v3'), ('x', '91'), ('k5', 'v5'), ('k4', 'v4'), ('k1', 'v1')]
    >>> len(cache)
    5
    >>> 'x' in cache
    True
    >>> cache.replace('k4', 'v4_new')
    True
    >>> cache.replace('notexists', 'value')
    False
    >>> cache.items()
    [('k4', 'v4_new'), ('k3', 'v3'), ('x', '91'), ('k5', 'v5'), ('k1', 'v1')]
"""

from heapq import heappush, heappop, heapify
from threading import Lock
from time import time
from weakref import ref

class Node(object):
    __slots__ = 'prev', 'next', 'expire', 'key', 'value', '__weakref__'

def normalize(key, value="", expire=None):
    if isinstance(key, tuple): hash_value, key = key
    elif not isinstance(key, str): raise ValueError('Key must be tuple or string. Got: %s' % key.__class__.__name__)
    if not isinstance(value, str): raise ValueError('Value must be string. Got: %s' % value.__class__.__name__)
    if len(value) > 1024 * 1024: raise ValueError('Value size too big: %d' % len(value))
    if expire is not None: expire = time() + expire
    return key, value, expire

class Memcache(object):
    def __init__(self, max_data_size):
        self.dict = {}
        list = self.list = Node()
        list.prev = list.next = list.expire = list.key = list.value = None
        list.prev = list.next = list
        self.lock = Lock()
        self.heap = []
        self.data_size = 0
        self.max_data_size = max_data_size
    def __len__(self):
        return len(self.dict)
    def __contains__(self, key):
        return key in self.dict
    def __iter__(self):
        return iter(self.items())
    def items(self):
        now = time()
        result = []
        append = result.append
        list = self.list
        self.lock.acquire()
        try:
            node = list.prev
            while node is not list:
                expire = node.expire
                if expire is not None and expire <= now: self._delete_node(node)
                else: append((node.key, node.value))
                node = node.prev
        finally: self.lock.release()
        return result
    def _delete_node(self, node):
        prev, next = node.prev, node.next
        prev.next = next
        next.prev = prev
        del self.dict[node.key]
        self.data_size -= (len(node.key) + len(node.value))
    def _find_node(self, key):
        now = time()
        node = self.dict.get(key)
        if node is None: return None
        expire = node.expire
        success = expire is None or expire > now
        prev, next = node.prev, node.next
        prev.next = next
        next.prev = prev
        return node
    def _create_node(self, key):
        self.dict[key] = node = Node()
        node.key = key
        node.value = None
        return node
    def _place_on_top(self, node):
        list = self.list
        node.next = list
        node.prev = prev_top = list.prev
        list.prev = prev_top.next = node
    def _set_node_value(self, node, value, expire):
        prev_value = node.value
        node.value, node.expire = value, expire
        if prev_value is None:
            self.data_size += len(node.key)
            self.data_size += len(value)
        else:
            self.data_size -= len(prev_value)
            self.data_size += len(value)
        if expire is not None: heappush(self.heap, (expire, ref(node)))
        self._delete_expired_nodes()
        self._conform_to_limits()
        if len(self.heap) > len(self.dict) * 2: self._pack_heap()
    def _delete_expired_nodes(self):
        now = time()
        heap = self.heap
        while heap:
            expire, node_ref = heap[0]
            if expire > now: break
            heappop(heap)
            node = node_ref()
            if node is not None and expire == node.expire: self._delete_node(node)
    def _conform_to_limits(self):
        list = self.list
        while self.data_size > self.max_data_size:
            bottom = list.next
            self._delete_node(bottom)
    def _pack_heap(self):
        new_heap = []
        for item in self.heap:
            expire, node_ref = item
            node = node_ref()
            if node is not None and expire == node.expire: new_heap.append(item)
        heapify(new_heap)
        self.heap = new_heap
    def get(self, key):
        key, _, _ = normalize(key)
        self.lock.acquire()
        try:
            node = self._find_node(key)
            if node is None: return None
            self._place_on_top(node)
            return node.value
        finally: self.lock.release()
    def set(self, key, value, expire=None):
        key, value, expire = normalize(key, value, expire)
        self.lock.acquire()
        try:
            node = self._find_node(key)
            if node is None: node = self._create_node(key)
            self._place_on_top(node)
            self._set_node_value(node, value, expire)
        finally: self.lock.release()
        return True
    def add(self, key, value, expire=None):
        key, value, expire = normalize(key, value, expire)
        self.lock.acquire()
        try:
            node = self._find_node(key)
            if node:
                self._place_on_top(node)
                return False
            node = self._create_node(key)
            self._place_on_top(node)
            self._set_node_value(node, value, expire)
        finally: self.lock.release()
        return True
    def replace(self, key, value, expire=None):
        key, value, expire = normalize(key, value, expire)
        self.lock.acquire()
        try:
            node = self._find_node(key)
            if node is None: return False
            self._place_on_top(node)
            self._set_node_value(node, value, expire)
        finally: self.lock.release()
        return True
    def delete(self, key):
        key, _, _ = normalize(key)
        node = self._find_node(key)
        if node is None: return 1
        self._delete_node(node)
        return 2
    def incr(self, key, delta=1):
        key, _, _ = normalize(key)
        self.lock.acquire()
        try:
            node = self._find_node(key)
            if node is None: return None
            self._place_on_top(node)
            try: value = int(node.value) + delta
            except ValueError: return None
            node.value = str(value)
        finally: self.lock.release()
        return value
    def decr(self, key, delta=1):
        key, _, _ = normalize(key)
        self.lock.acquire()
        try:
            node = self._find_node(key)
            if node is None: return None
            self._place_on_top(node)
            try: value = int(node.value) - delta
            except ValueError: return None
            node.value = str(value)
        finally: self.lock.release()
        return value
    def flush_all(self):
        self.lock.acquire()
        try:
            self.dict.clear()
            self.heap = []
            self.list.prev = self.list.next = self.list
            for node in self.dict.itervalues():
                node.prev = node.next = None
        finally: self.lock.release()
<|endoftext|>