<commit_before>import os
import json
import logging
import lzma
from binascii import hexlify
from io import BytesIO
from urllib.parse import urlparse
from uuid import UUID
from . import engine as UnityEngine
from .enums import RuntimePlatform
from .utils import BinaryReader

__author__ = "Jerome Leclanche"
__email__ = "jerome@leclan.ch"
__version__ = "0.1"

SIGNATURE_WEB = "UnityWeb"
SIGNATURE_RAW = "UnityRaw"

def get_asset(name):
	return os.path.join(os.path.dirname(__file__), name)

with open(get_asset("strings.dat"), "rb") as f:
	STRINGS_DAT = f.read()

with open(get_asset("classes.json"), "r") as f:
	UNITY_CLASSES = json.load(f)

def UnityClass(i):
	return UNITY_CLASSES.get(str(i), "<Unknown #%i>" % (i))

def load_object(clsname, obj):
	if hasattr(UnityEngine, clsname):
		obj = getattr(UnityEngine, clsname)(obj)

	return obj

class TypeTree:
	NULL = "(null)"

	def __init__(self):
		self.children = []
		self.version = 0
		self.is_array = False
		self.size = 0
		self.index = 0
		self.flags = 0
		self.type = self.NULL
		self.name = self.NULL

	def __repr__(self):
		return "<%s %s (size=%r, index=%r, is_array=%r, flags=%r)>" % (
			self.type, self.name, self.size, self.index, self.is_array, self.flags
		)

	@property
	def post_align(self):
		return self.flags & 0x4000

	def load_blob(self, buf):
		self.num_nodes = buf.read_uint()
		self.buffer_bytes = buf.read_uint()
		node_data = BytesIO(buf.read(24 * self.num_nodes))
		self.data = buf.read(self.buffer_bytes)

		parents = [self]

		buf = BinaryReader(node_data)

		for i in range(self.num_nodes):
			version = buf.read_int16()
			depth = buf.read_byte()
			is_array = buf.read_byte()
			type_offset = buf.read_int()
			name_offset = buf.read_int()
			size = buf.read_int()
			index = buf.read_uint()
			flags = buf.read_int()

			type = self.get_string(type_offset)
			name = self.get_string(name_offset)

			if depth == 0:
				curr = self
			else:
				while len(parents) > depth:
					parents.pop()
				curr = TypeTree()
				parents[-1].children.append(curr)
				parents.append(curr)

			curr.type = type
			curr.name = name
			curr.size = size
			curr.index = index
			curr.is_array = is_array
			curr.version = version
			curr.flags = flags

	def get_string(self, offset):
		if offset < 0:
			offset &= 0x7fffffff
			data = STRINGS_DAT
		elif offset < self.buffer_bytes:
			data = self.data
		else:
			return self.NULL
		return data[offset:].partition(b"\0")[0].decode("utf-8")

class TypeMetadata:
	default_instance = None

	@classmethod
	def default(cls):
		if not cls.default_instance:
			cls.default_instance = cls()
			with open(get_asset("structs.dat"), "rb") as f:
				cls.default_instance.load(BinaryReader(f))
		return cls.default_instance

	def __init__(self):
		self.type_trees = {}
		self.hashes = {}

	def load(self, buf):
		offset = buf.tell()
		self.generator_version = buf.read_string()
		self.target_platform = RuntimePlatform(buf.read_uint())

		# if format >= 13
		self.has_type_trees = buf.read_boolean()
		self.num_types = buf.read_int()

		for i in range(self.num_types):
			class_id = buf.read_int()
			if class_id < 0:
				hash = buf.read(0x20)
			else:
				hash = buf.read(0x10)

			self.hashes[class_id] = hash

			if self.has_type_trees:
				tree = TypeTree()
				tree.load_blob(buf)
				self.type_trees[class_id] = tree

class ObjectInfo:
	def __init__(self, asset):
		self.asset = asset

	def __repr__(self):
		return "<%s %i>" % (self.type, self.class_id)

	@property
	def type(self):
		if self.type_id > 0:
			return UnityClass(self.type_id)
		else:
			return "<N/A>"

	def load(self, buf):
		self.path_id = self.asset.read_id(buf)
		self.data_offset = buf.read_uint() + self.asset.data_offset
		self.size = buf.read_uint()
		self.type_id = buf.read_int()
		self.class_id = buf.read_int16()

		if self.asset.format <= 10:
			self.is_destroyed = bool(buf.read_int16())
		elif self.asset.format >= 11:
			self.unk0 = buf.read_int16()

			if self.asset.format >= 15:
				self.unk1 = buf.read_byte()

	def read(self):
		type = self.asset.types[self.type_id]
		buf = BinaryReader(self.asset.data)
		buf.seek(self.data_offset)
		return self.read_value(type, buf)

	def read_value(self, type, buf):
		align = False
		t = type.type
		first_child = type.children[0] if type.children else TypeTree()
		if t == "bool":
			result = buf.read_boolean()
		elif t == "UInt8":
			result = buf.read_byte()
		elif t == "SInt16":
			result = buf.read_int16()
		elif t == "UInt16":
			result = buf.read_uint16()
		elif t == "UInt64":
			result = buf.read_int64()
		elif t == "unsigned int":
			result = buf.read_uint()
		elif t == "int":
			result = buf.read_int()
		elif t == "float":
			buf.align()
			result = buf.read_float()
		elif t == "string":
			size = buf.read_uint()
			result = buf.read_string(size)
			align = type.children[0].post_align
		else:
			if type.is_array:
				first_child = type

			if t.startswith("PPtr<"):
				result = ObjectPointer(self.asset)
				result.load(buf)
				if not result:
					result = None

			elif first_child and first_child.is_array:
				align = first_child.post_align
				size = buf.read_uint()
				array_type = first_child.children[1]
				if array_type.type in ("char", "UInt8"):
					result = buf.read(size)
				else:
					result = []
					for i in range(size):
						result.append(self.read_value(array_type, buf))
			else:
				result = {}

				for child in type.children:
					result[child.name] = self.read_value(child, buf)

				result = load_object(t, result)
				if t == "StreamedResource":
					if self.asset.bundle:
						result.asset = self.asset.bundle.get_asset(result.source)
					else:
						logging.warning("StreamedResource not available without bundle")
						result.asset = None

		if align or type.post_align:
			buf.align()

		return result

class ObjectPointer:
	def __init__(self, asset):
		self.source_asset = asset

	def __bool__(self):
		return not (self.file_id == 0 and self.path_id == 0)

	def load(self, buf):
		self.file_id = buf.read_int()
		self.path_id = self.source_asset.read_id(buf)

	def __repr__(self):
		return "%s(file_id=%r, path_id=%r)" % (
			self.__class__.__name__, self.file_id, self.path_id
		)

	@property
	def asset(self):
		return self.source_asset.asset_refs[self.file_id]

	def resolve(self):
		return self.asset.objects[self.path_id]

class Asset:
	@classmethod
	def from_bundle(cls, bundle, buf):
		ret = cls()
		offset = buf.tell()
		if not bundle.compressed:
			ret.name = buf.read_string()
			header_size = buf.read_uint()
			size = buf.read_uint()
		else:
			header_size = bundle.asset_header_size

		# FIXME: this offset needs to be explored more
		ofs = buf.tell()
		if bundle.compressed:
			dec = lzma.LZMADecompressor()
			data = dec.decompress(buf.read())
			data = BytesIO(data[header_size:])
		else:
			if ret.is_resource:
				buf.seek(offset + header_size - 4 - len(ret.name))
			else:
				buf.seek(offset + header_size - 4)
			data = BytesIO(buf.read())
		ret.data = BinaryReader(data, endian=">")
		buf.seek(ofs)
		ret.bundle = bundle
		return ret

	@classmethod
	def from_file(cls, file):
		ret = cls()
		ret.name = file.name
		ret.data = BinaryReader(BytesIO(file.read()), endian=">")
		ret.load(ret.data)
		return ret

	def __init__(self):
		self.objects = {}
		self.adds = []
		self.asset_refs = [self]
		self.types = {}
		self.bundle = None
		self.name = ""

	def __repr__(self):
		return "<%s %s>" % (self.__class__.__name__, self.name)

	@property
	def is_resource(self):
		return self.name.endswith(".resource")

	def load(self, buf):
		self.metadata_size = buf.read_uint()
		self.file_size = buf.read_uint()
		self.format = buf.read_uint()
		self.data_offset = buf.read_uint()

		if self.format >= 9:
			self.endianness = buf.read_uint()
			if self.endianness == 0:
				buf.endian = "<"

		self.tree = TypeMetadata()
		self.tree.load(buf)

		self.num_objects = buf.read_uint()
		for i in range(self.num_objects):
			if self.format >= 14:
				buf.align()
			obj = ObjectInfo(self)
			obj.load(buf)
			self.register_object(obj)

		if self.format >= 11:
			num_adds = buf.read_uint()
			for i in range(num_adds):
				if self.format >= 14:
					buf.align()
				id = self.read_id(buf)
				self.adds.append((id, buf.read_int()))

		if self.format >= 14:
			num_refs = buf.read_uint()
			for i in range(num_refs):
				ref = AssetRef()
				ref.load(buf)
				self.asset_refs.append(ref)

		unk_string = buf.read_string()
		assert not unk_string, unk_string

	def read_id(self, buf):
		if self.format >= 14:
			return buf.read_int64()
		else:
			return buf.read_int()

	def register_object(self, obj):
		if obj.type_id in self.tree.type_trees:
			self.types[obj.type_id] = self.tree.type_trees[obj.type_id]
		elif obj.type_id not in self.types:
			tree = TypeMetadata.default().type_trees
			if obj.class_id in tree:
				self.types[obj.type_id] = TypeMetadata.default().type_trees[obj.class_id]
			else:
				logging.warning("%r absent from structs.dat", obj.class_id)
				self.types[obj.type_id] = None

		if obj.path_id in self.objects:
			raise ValueError("Duplicate asset object: %r" % (obj))

		self.objects[obj.path_id] = obj

class AssetRef:
	def load(self, buf):
		self.asset_path = buf.read_string()
		self.guid = UUID(hexlify(buf.read(16)).decode("utf-8"))
		self.type = buf.read_int()
		self.file_path = buf.read_string()
		self.asset = None  # TODO loadrefs

	def __repr__(self):
		return "<%s (asset_path=%r, guid=%r, type=%r, file_path=%r)>" % (
			self.__class__.__name__, self.asset_path, self.guid, self.type, self.file_path
		)

class AssetBundle:
	def __init__(self):
		self.assets = []

	@property
	def compressed(self):
		return self.signature == SIGNATURE_WEB

	def load(self, file):
		buf = BinaryReader(file, endian=">")

		self.signature = buf.read_string()
		self.format_version = buf.read_int()
		self.unity_version = buf.read_string()
		self.generator_version = buf.read_string()
		self.file_size = buf.read_uint()
		self.header_size = buf.read_int()

		self.file_count = buf.read_int()
		self.bundle_count = buf.read_int()

		if self.format_version >= 2:
			self.bundle_size = buf.read_uint()  # without header_size

			if self.format_version >= 3:
				self.uncompressed_bundle_size = buf.read_uint()  # without header_size

		if self.header_size >= 60:
			self.compressed_file_size = buf.read_uint()  # with header_size
			self.asset_header_size = buf.read_uint()

		assert self.signature in (SIGNATURE_RAW, SIGNATURE_WEB)

		# Preload assets
		buf.seek(self.header_size)
		if not self.compressed:
			self.num_assets = buf.read_int()
		else:
			self.num_assets = 1
		for i in range(self.num_assets):
			asset = Asset.from_bundle(self, buf)
			if not asset.is_resource:
				asset.load(asset.data)
			self.assets.append(asset)

	def get_asset(self, url):
		if not url:
			return None
		u = urlparse(url)
		assert u.scheme == "archive"

		archive, name = os.path.split(u.path)
		# Ignore the archive for now

		for asset in self.assets:
			if asset.name == name:
				return asset
		raise KeyError("No such asset: %r" % (name))

def load(file):
	ret = AssetBundle()
	ret.load(file)
	return ret
<commit_msg>Simplify TypeTree loading<commit_after>import os
import json
import logging
import lzma
from binascii import hexlify
from io import BytesIO
from urllib.parse import urlparse
from uuid import UUID
from . import engine as UnityEngine
from .enums import RuntimePlatform
from .utils import BinaryReader

__author__ = "Jerome Leclanche"
__email__ = "jerome@leclan.ch"
__version__ = "0.1"

SIGNATURE_WEB = "UnityWeb"
SIGNATURE_RAW = "UnityRaw"

def get_asset(name):
	return os.path.join(os.path.dirname(__file__), name)

with open(get_asset("strings.dat"), "rb") as f:
	STRINGS_DAT = f.read()

with open(get_asset("classes.json"), "r") as f:
	UNITY_CLASSES = json.load(f)

def UnityClass(i):
	return UNITY_CLASSES.get(str(i), "<Unknown #%i>" % (i))

def load_object(clsname, obj):
	if hasattr(UnityEngine, clsname):
		obj = getattr(UnityEngine, clsname)(obj)

	return obj

class TypeTree:
	NULL = "(null)"

	def __init__(self):
		self.children = []
		self.version = 0
		self.is_array = False
		self.size = 0
		self.index = 0
		self.flags = 0
		self.type = self.NULL
		self.name = self.NULL

	def __repr__(self):
		return "<%s %s (size=%r, index=%r, is_array=%r, flags=%r)>" % (
			self.type, self.name, self.size, self.index, self.is_array, self.flags
		)

	@property
	def post_align(self):
		return self.flags & 0x4000

	def load_blob(self, buf):
		num_nodes = buf.read_uint()
		self.buffer_bytes = buf.read_uint()
		node_data = BytesIO(buf.read(24 * num_nodes))
		self.data = buf.read(self.buffer_bytes)

		parents = [self]

		buf = BinaryReader(node_data)

		for i in range(num_nodes):
			version = buf.read_int16()
			depth = buf.read_byte()

			if depth == 0:
				curr = self
			else:
				while len(parents) > depth:
					parents.pop()
				curr = TypeTree()
				parents[-1].children.append(curr)
				parents.append(curr)

			curr.version = version
			curr.is_array = buf.read_byte()
			curr.type = self.get_string(buf.read_int())
			curr.name = self.get_string(buf.read_int())
			curr.size = buf.read_int()
			curr.index = buf.read_uint()
			curr.flags = buf.read_int()

	def get_string(self, offset):
		if offset < 0:
			offset &= 0x7fffffff
			data = STRINGS_DAT
		elif offset < self.buffer_bytes:
			data = self.data
		else:
			return self.NULL
		return data[offset:].partition(b"\0")[0].decode("utf-8")

class TypeMetadata:
	default_instance = None

	@classmethod
	def default(cls):
		if not cls.default_instance:
			cls.default_instance = cls()
			with open(get_asset("structs.dat"), "rb") as f:
				cls.default_instance.load(BinaryReader(f))
		return cls.default_instance

	def __init__(self):
		self.type_trees = {}
		self.hashes = {}

	def load(self, buf):
		offset = buf.tell()
		self.generator_version = buf.read_string()
		self.target_platform = RuntimePlatform(buf.read_uint())

		# if format >= 13
		self.has_type_trees = buf.read_boolean()
		self.num_types = buf.read_int()

		for i in range(self.num_types):
			class_id = buf.read_int()
			if class_id < 0:
				hash = buf.read(0x20)
			else:
				hash = buf.read(0x10)

			self.hashes[class_id] = hash

			if self.has_type_trees:
				tree = TypeTree()
				tree.load_blob(buf)
				self.type_trees[class_id] = tree

class ObjectInfo:
	def __init__(self, asset):
		self.asset = asset

	def __repr__(self):
		return "<%s %i>" % (self.type, self.class_id)

	@property
	def type(self):
		if self.type_id > 0:
			return UnityClass(self.type_id)
		else:
			return "<N/A>"

	def load(self, buf):
		self.path_id = self.asset.read_id(buf)
		self.data_offset = buf.read_uint() + self.asset.data_offset
		self.size = buf.read_uint()
		self.type_id = buf.read_int()
		self.class_id = buf.read_int16()

		if self.asset.format <= 10:
			self.is_destroyed = bool(buf.read_int16())
		elif self.asset.format >= 11:
			self.unk0 = buf.read_int16()

			if self.asset.format >= 15:
				self.unk1 = buf.read_byte()

	def read(self):
		type = self.asset.types[self.type_id]
		buf = BinaryReader(self.asset.data)
		buf.seek(self.data_offset)
		return self.read_value(type, buf)

	def read_value(self, type, buf):
		align = False
		t = type.type
		first_child = type.children[0] if type.children else TypeTree()
		if t == "bool":
			result = buf.read_boolean()
		elif t == "UInt8":
			result = buf.read_byte()
		elif t == "SInt16":
			result = buf.read_int16()
		elif t == "UInt16":
			result = buf.read_uint16()
		elif t == "UInt64":
			result = buf.read_int64()
		elif t == "unsigned int":
			result = buf.read_uint()
		elif t == "int":
			result = buf.read_int()
		elif t == "float":
			buf.align()
			result = buf.read_float()
		elif t == "string":
			size = buf.read_uint()
			result = buf.read_string(size)
			align = type.children[0].post_align
		else:
			if type.is_array:
				first_child = type

			if t.startswith("PPtr<"):
				result = ObjectPointer(self.asset)
				result.load(buf)
				if not result:
					result = None

			elif first_child and first_child.is_array:
				align = first_child.post_align
				size = buf.read_uint()
				array_type = first_child.children[1]
				if array_type.type in ("char", "UInt8"):
					result = buf.read(size)
				else:
					result = []
					for i in range(size):
						result.append(self.read_value(array_type, buf))
			else:
				result = {}

				for child in type.children:
					result[child.name] = self.read_value(child, buf)

				result = load_object(t, result)
				if t == "StreamedResource":
					if self.asset.bundle:
						result.asset = self.asset.bundle.get_asset(result.source)
					else:
						logging.warning("StreamedResource not available without bundle")
						result.asset = None

		if align or type.post_align:
			buf.align()

		return result

class ObjectPointer:
	def __init__(self, asset):
		self.source_asset = asset

	def __bool__(self):
		return not (self.file_id == 0 and self.path_id == 0)

	def load(self, buf):
		self.file_id = buf.read_int()
		self.path_id = self.source_asset.read_id(buf)

	def __repr__(self):
		return "%s(file_id=%r, path_id=%r)" % (
			self.__class__.__name__, self.file_id, self.path_id
		)

	@property
	def asset(self):
		return self.source_asset.asset_refs[self.file_id]

	def resolve(self):
		return self.asset.objects[self.path_id]

class Asset:
	@classmethod
	def from_bundle(cls, bundle, buf):
		ret = cls()
		offset = buf.tell()
		if not bundle.compressed:
			ret.name = buf.read_string()
			header_size = buf.read_uint()
			size = buf.read_uint()
		else:
			header_size = bundle.asset_header_size

		# FIXME: this offset needs to be explored more
		ofs = buf.tell()
		if bundle.compressed:
			dec = lzma.LZMADecompressor()
			data = dec.decompress(buf.read())
			data = BytesIO(data[header_size:])
		else:
			if ret.is_resource:
				buf.seek(offset + header_size - 4 - len(ret.name))
			else:
				buf.seek(offset + header_size - 4)
			data = BytesIO(buf.read())
		ret.data = BinaryReader(data, endian=">")
		buf.seek(ofs)
		ret.bundle = bundle
		return ret

	@classmethod
	def from_file(cls, file):
		ret = cls()
		ret.name = file.name
		ret.data = BinaryReader(BytesIO(file.read()), endian=">")
		ret.load(ret.data)
		return ret

	def __init__(self):
		self.objects = {}
		self.adds = []
		self.asset_refs = [self]
		self.types = {}
		self.bundle = None
		self.name = ""

	def __repr__(self):
		return "<%s %s>" % (self.__class__.__name__, self.name)

	@property
	def is_resource(self):
		return self.name.endswith(".resource")

	def load(self, buf):
		self.metadata_size = buf.read_uint()
		self.file_size = buf.read_uint()
		self.format = buf.read_uint()
		self.data_offset = buf.read_uint()

		if self.format >= 9:
			self.endianness = buf.read_uint()
			if self.endianness == 0:
				buf.endian = "<"

		self.tree = TypeMetadata()
		self.tree.load(buf)

		self.num_objects = buf.read_uint()
		for i in range(self.num_objects):
			if self.format >= 14:
				buf.align()
			obj = ObjectInfo(self)
			obj.load(buf)
			self.register_object(obj)

		if self.format >= 11:
			num_adds = buf.read_uint()
			for i in range(num_adds):
				if self.format >= 14:
					buf.align()
				id = self.read_id(buf)
				self.adds.append((id, buf.read_int()))

		if self.format >= 14:
			num_refs = buf.read_uint()
			for i in range(num_refs):
				ref = AssetRef()
				ref.load(buf)
				self.asset_refs.append(ref)

		unk_string = buf.read_string()
		assert not unk_string, unk_string

	def read_id(self, buf):
		if self.format >= 14:
			return buf.read_int64()
		else:
			return buf.read_int()

	def register_object(self, obj):
		if obj.type_id in self.tree.type_trees:
			self.types[obj.type_id] = self.tree.type_trees[obj.type_id]
		elif obj.type_id not in self.types:
			tree = TypeMetadata.default().type_trees
			if obj.class_id in tree:
				self.types[obj.type_id] = TypeMetadata.default().type_trees[obj.class_id]
			else:
				logging.warning("%r absent from structs.dat", obj.class_id)
				self.types[obj.type_id] = None

		if obj.path_id in self.objects:
			raise ValueError("Duplicate asset object: %r" % (obj))

		self.objects[obj.path_id] = obj

class AssetRef:
	def load(self, buf):
		self.asset_path = buf.read_string()
		self.guid = UUID(hexlify(buf.read(16)).decode("utf-8"))
		self.type = buf.read_int()
		self.file_path = buf.read_string()
		self.asset = None  # TODO loadrefs

	def __repr__(self):
		return "<%s (asset_path=%r, guid=%r, type=%r, file_path=%r)>" % (
			self.__class__.__name__, self.asset_path, self.guid, self.type, self.file_path
		)

class AssetBundle:
	def __init__(self):
		self.assets = []

	@property
	def compressed(self):
		return self.signature == SIGNATURE_WEB

	def load(self, file):
		buf = BinaryReader(file, endian=">")

		self.signature = buf.read_string()
		self.format_version = buf.read_int()
		self.unity_version = buf.read_string()
		self.generator_version = buf.read_string()
		self.file_size = buf.read_uint()
		self.header_size = buf.read_int()

		self.file_count = buf.read_int()
		self.bundle_count = buf.read_int()

		if self.format_version >= 2:
			self.bundle_size = buf.read_uint()  # without header_size

			if self.format_version >= 3:
				self.uncompressed_bundle_size = buf.read_uint()  # without header_size

		if self.header_size >= 60:
			self.compressed_file_size = buf.read_uint()  # with header_size
			self.asset_header_size = buf.read_uint()

		assert self.signature in (SIGNATURE_RAW, SIGNATURE_WEB)

		# Preload assets
		buf.seek(self.header_size)
		if not self.compressed:
			self.num_assets = buf.read_int()
		else:
			self.num_assets = 1
		for i in range(self.num_assets):
			asset = Asset.from_bundle(self, buf)
			if not asset.is_resource:
				asset.load(asset.data)
			self.assets.append(asset)

	def get_asset(self, url):
		if not url:
			return None
		u = urlparse(url)
		assert u.scheme == "archive"

		archive, name = os.path.split(u.path)
		# Ignore the archive for now

		for asset in self.assets:
			if asset.name == name:
				return asset
		raise KeyError("No such asset: %r" % (name))

def load(file):
	ret = AssetBundle()
	ret.load(file)
	return ret
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding=utf-8
import time
import os
import importlib
from cyclone.web import RequestHandler
from cyclone.websocket import WebSocketHandler
from toughlib import dispatch,logger

class Permit():
    """ 权限菜单管理
    @TODO: ⚠注意，这个权限管理在多进程时可能会有数据一致性问题，不过现在展示不解决它
    """
    routes = {}
    handlers = {}

    def add_route(self, handle_cls, path, name, category, handle_params={}, is_menu=False, order=time.time(),
                  is_open=True, oem=False):
        """ 注册权限
        """
        if not path: return
        if path in self.routes:
            if self.routes[path].get('oem'):
                return

        self.routes[path] = dict(
            path=path,  # 权限url路径
            name=name,  # 权限名称
            category=category,  # 权限目录
            is_menu=is_menu,  # 是否在边栏显示为菜单
            oprs=[],  # 关联的操作员
            order=order,  # 排序
            is_open=is_open,  # 是否开放授权
            oem=oem #是否定制功能
        )
        self.add_handler(handle_cls, path, handle_params)

    def add_handler(self, handle_cls, path, handle_params={}):
        print ("add handler [%s:%s]" % (path, repr(handle_cls)))
        self.handlers[path]= (path, handle_cls, handle_params)

    @property
    def all_handlers(self):
        return self.handlers.values()

    def get_route(self, path):
        """ 获取一个权限资源
        """
        return self.routes.get(path)

    def bind_super(self, opr):
        """ 为超级管理员授权所有权限
        """
        for path in self.routes:
            route = self.routes.get(path)
            route['oprs'].append(opr)

    def bind_opr(self, opr, path):
        """ 为操作员授权
        """
        if not path or path not in self.routes:
            return
        oprs = self.routes[path]['oprs']
        if opr not in oprs:
            oprs.append(opr)

    def unbind_opr(self, opr, path=None):
        """ 接触操作员与权限关联
        """
        if path:
            self.routes[path]['oprs'].remove(opr)
        else:
            for path in self.routes:
                route = self.routes.get(path)
                if route and opr in route['oprs']:
                    route['oprs'].remove(opr)

    def check_open(self, path):
        """ 检查权限是否开放授权
        """
        route = self.routes[path]
        return 'is_open' in route and route['is_open']

    def check_opr_category(self, opr, category):
        """ 检查权限是否在指定目录下
        """
        for path in self.routes:
            route = self.routes[path]
            if opr in route['oprs'] and route['category'] == category:
                return True
        return False

    def build_menus(self, order_cats=[]):
        """ 生成全局内存菜单"""
        menus = [{'category': _cat, 'items': []} for _cat in order_cats]
        for path in self.routes:
            route = self.routes[path]
            for menu in menus:
                if route['category'] == menu['category']:
                    menu['items'].append(route)
        return menus

    def match(self, opr, path):
        """ 检查操作员是否匹配资源
        """
        if not path or not opr:
            return False
        if path not in self.routes:
            return False
        return opr in self.routes[path]['oprs']

    def route(self, url_pattern, menuname=None, category=None, is_menu=False, order=0, is_open=True,oem=False):
        selfobj = self

        def handler_wapper(cls):
            assert (issubclass(cls, RequestHandler) or issubclass(cls, WebSocketHandler))
            if not menuname:
                self.add_handler(cls, url_pattern)
            else:
                selfobj.add_route(cls, url_pattern, menuname, category,
                        order=order, is_menu=is_menu, is_open=is_open,oem=oem)
            return cls

        return handler_wapper

# 全局实例
permit = Permit()

def load_handlers(handler_path=None, pkg_prefix=None, excludes=[]):
    _excludes = ['__init__', 'base', '.svn', '.DS_Store', 'views'] + excludes
    hds = set(os.path.splitext(it)[0] for it in os.listdir(handler_path))
    hds = [it for it in hds if it not in _excludes]
    for hd in hds:
        try:
            sub_module = os.path.join(handler_path, hd)
            if os.path.isdir(sub_module):
                print ('load sub module %s' % hd)
                load_handlers(
                    handler_path=sub_module,
                    pkg_prefix="{0}.{1}".format(pkg_prefix, hd)
                )

            _hd = "{0}.{1}".format(pkg_prefix, hd)
            print ('load_module %s' % _hd)
            importlib.import_module(_hd)
        except Exception as err:
            logger.error("%s, skip module %s.%s" % (str(err),pkg_prefix,hd))
            import traceback
            traceback.print_exc()
            continue

def load_events(event_path=None,pkg_prefix=None,excludes=[],event_params={}):
    _excludes = ['__init__','settings'] + excludes
    evs = set(os.path.splitext(it)[0] for it in os.listdir(event_path))
    evs = [it for it in evs if it not in _excludes]
    for ev in evs:
        try:
            sub_module = os.path.join(event_path, ev)
            if os.path.isdir(sub_module):
                print ('load sub event %s' % ev)
                load_events(
                    event_path=sub_module,
                    pkg_prefix="{0}.{1}".format(pkg_prefix, ev),
                    excludes=excludes,
                    event_params=event_params,
                )
            _ev = "{0}.{1}".format(pkg_prefix, ev)
            print ('load_event %s with params:%s' % (_ev,repr(event_params)))
            dispatch.register(importlib.import_module(_ev).__call__(**event_params))
        except Exception as err:
            logger.error("%s, skip module %s.%s" % (str(err),pkg_prefix,ev))
            import traceback
            traceback.print_exc()
            continue

<commit_msg>event params<commit_after>#!/usr/bin/env python
# coding=utf-8
import time
import os
import importlib
from cyclone.web import RequestHandler
from cyclone.websocket import WebSocketHandler
from toughlib import dispatch,logger

class Permit():
    """ 权限菜单管理
    @TODO: ⚠注意，这个权限管理在多进程时可能会有数据一致性问题，不过现在展示不解决它
    """
    routes = {}
    handlers = {}

    def add_route(self, handle_cls, path, name, category, handle_params={}, is_menu=False, order=time.time(),
                  is_open=True, oem=False):
        """ 注册权限
        """
        if not path: return
        if path in self.routes:
            if self.routes[path].get('oem'):
                return

        self.routes[path] = dict(
            path=path,  # 权限url路径
            name=name,  # 权限名称
            category=category,  # 权限目录
            is_menu=is_menu,  # 是否在边栏显示为菜单
            oprs=[],  # 关联的操作员
            order=order,  # 排序
            is_open=is_open,  # 是否开放授权
            oem=oem #是否定制功能
        )
        self.add_handler(handle_cls, path, handle_params)

    def add_handler(self, handle_cls, path, handle_params={}):
        logger.info("add handler [%s:%s]" % (path, repr(handle_cls)))
        self.handlers[path]= (path, handle_cls, handle_params)

    @property
    def all_handlers(self):
        return self.handlers.values()

    def get_route(self, path):
        """ 获取一个权限资源
        """
        return self.routes.get(path)

    def bind_super(self, opr):
        """ 为超级管理员授权所有权限
        """
        for path in self.routes:
            route = self.routes.get(path)
            route['oprs'].append(opr)

    def bind_opr(self, opr, path):
        """ 为操作员授权
        """
        if not path or path not in self.routes:
            return
        oprs = self.routes[path]['oprs']
        if opr not in oprs:
            oprs.append(opr)

    def unbind_opr(self, opr, path=None):
        """ 接触操作员与权限关联
        """
        if path:
            self.routes[path]['oprs'].remove(opr)
        else:
            for path in self.routes:
                route = self.routes.get(path)
                if route and opr in route['oprs']:
                    route['oprs'].remove(opr)

    def check_open(self, path):
        """ 检查权限是否开放授权
        """
        route = self.routes[path]
        return 'is_open' in route and route['is_open']

    def check_opr_category(self, opr, category):
        """ 检查权限是否在指定目录下
        """
        for path in self.routes:
            route = self.routes[path]
            if opr in route['oprs'] and route['category'] == category:
                return True
        return False

    def build_menus(self, order_cats=[]):
        """ 生成全局内存菜单"""
        menus = [{'category': _cat, 'items': []} for _cat in order_cats]
        for path in self.routes:
            route = self.routes[path]
            for menu in menus:
                if route['category'] == menu['category']:
                    menu['items'].append(route)
        return menus

    def match(self, opr, path):
        """ 检查操作员是否匹配资源
        """
        if not path or not opr:
            return False
        if path not in self.routes:
            return False
        return opr in self.routes[path]['oprs']

    def route(self, url_pattern, menuname=None, category=None, is_menu=False, order=0, is_open=True,oem=False):
        selfobj = self

        def handler_wapper(cls):
            assert (issubclass(cls, RequestHandler) or issubclass(cls, WebSocketHandler))
            if not menuname:
                self.add_handler(cls, url_pattern)
            else:
                selfobj.add_route(cls, url_pattern, menuname, category,
                        order=order, is_menu=is_menu, is_open=is_open,oem=oem)
            return cls

        return handler_wapper

# 全局实例
permit = Permit()

def load_handlers(handler_path=None, pkg_prefix=None, excludes=[]):
    _excludes = ['__init__', 'base', '.svn', '.DS_Store', 'views'] + excludes
    hds = set(os.path.splitext(it)[0] for it in os.listdir(handler_path))
    hds = [it for it in hds if it not in _excludes]
    for hd in hds:
        try:
            sub_module = os.path.join(handler_path, hd)
            if os.path.isdir(sub_module):
                logger.info('load sub module %s' % hd)
                load_handlers(
                    handler_path=sub_module,
                    pkg_prefix="{0}.{1}".format(pkg_prefix, hd)
                )

            _hd = "{0}.{1}".format(pkg_prefix, hd)
            logger.info('load_module %s' % _hd)
            importlib.import_module(_hd)
        except Exception as err:
            logger.error("%s, skip module %s.%s" % (str(err),pkg_prefix,hd))
            import traceback
            traceback.print_exc()
            continue

def load_events(event_path=None,pkg_prefix=None,excludes=[],event_params={}):
    _excludes = ['__init__','settings'] + excludes
    evs = set(os.path.splitext(it)[0] for it in os.listdir(event_path))
    evs = [it for it in evs if it not in _excludes]
    for ev in evs:
        try:
            sub_module = os.path.join(event_path, ev)
            if os.path.isdir(sub_module):
                logger.info('load sub event %s' % ev)
                load_events(
                    event_path=sub_module,
                    pkg_prefix="{0}.{1}".format(pkg_prefix, ev),
                    excludes=excludes,
                    event_params=event_params,
                )
            _ev = "{0}.{1}".format(pkg_prefix, ev)
            logger.info('load_event %s with params:%s' % (_ev,repr(event_params)))
            dispatch.register(importlib.import_module(_ev).__call__(**event_params))
        except Exception as err:
            logger.error("%s, skip module %s.%s" % (str(err),pkg_prefix,ev))
            import traceback
            traceback.print_exc()
            continue

<|endoftext|>


<commit_before>"""
This is an experiment in splitting up the monolithic model view class into
smaller, more reusable parts, and using class-based views at the same time
(not the generic class based views, though)
"""

import json

from django import forms
from django.contrib import messages
from django.core.exceptions import ImproperlyConfigured
from django.forms.models import modelform_factory, model_to_dict
from django.http import Http404, HttpResponse, HttpResponseRedirect
from django.shortcuts import get_object_or_404, redirect
from django.utils.translation import ugettext as _
from django.views.generic.base import TemplateView

from towel.forms import BatchForm, towel_formfield_callback
from towel.paginator import Paginator, EmptyPage, InvalidPage
from towel.utils import changed_regions, safe_queryset_and

class ModelResourceView(TemplateView):
    base_template = 'base.html'
    model = None
    queryset = None
    template_name_suffix = None

    def get_context_data(self, **kwargs):
        opts = self.model._meta
        context = {
            'base_template': self.base_template,
            'verbose_name': opts.verbose_name,
            'verbose_name_plural': opts.verbose_name_plural,
            'view': self,

            # TODO add_url
            # TODO list_url
            # TODO adding_allowed
            }
        context.update(kwargs)
        return context

    def get_template_names(self):
        opts = self.model._meta
        names = [
            '{}/{}{}.html'.format(opts.app_label, opts.module_name,
                self.template_name_suffix),
            'resources/object{}.html'.format(self.template_name_suffix),
            ]
        if self.template_name:
            names.insert(0, self.template_name)
        return names

    def get_queryset(self):
        if self.queryset is not None:
            return self.queryset._clone()
        elif self.model is not None:
            return self.model._default_manager.all()
        else:
            raise ImproperlyConfigured("'%s' must define 'queryset' or 'model'"
                                       % self.__class__.__name__)

    def get_object(self):
        return get_object_or_404(self.get_queryset(), **self.kwargs)

    CREATE = type('CREATE', (), {})
    UPDATE = type('UPDATE', (), {})
    DELETE = type('DELETE', (), {})

    def allow(self, action, object=None, **kwargs):
        if action in (self.CREATE, self.UPDATE):
            return True
        return False

    default_messages = {
        CREATE: (messages.SUCCESS,
            _('The new object has been successfully created.')),
        UPDATE: (messages.SUCCESS,
            _('The object has been successfully updated.')),
        DELETE: (messages.SUCCESS,
            _('The object has been successfully deleted.')),

        'adding_denied': (messages.ERROR,
            _('You are not allowed to add objects.')),
        'editing_denied': (messages.ERROR,
            _('You are not allowed to edit this object.')),
        'deletion_denied': (messages.ERROR,
            _('You are not allowed to delete this object.')),
        'deletion_denied_related': (messages.ERROR,
            _('Deletion not allowed: There are %(pretty_classes)s related to this object.')),
    }

    def add_message(self, message, level=None, **kwargs):
        if message in self.default_messages:
            level, message = self.default_messages[message]
        messages.add_message(self.request, level, message, **kwargs)

class ListView(ModelResourceView):
    paginate_by = None
    search_form = None
    template_name_suffix = '_list'

    def get_paginate_by(self, queryset):
        """
        if self.paginate_all_allowed and self.request.GET.get('all'):
            return None
        """
        return self.paginate_by

    def get_context_data(self, object_list, **kwargs):
        context = super(ListView, self).get_context_data(
            object_list=object_list, **kwargs)

        paginate_by = self.get_paginate_by(object_list)
        if paginate_by:
            paginator = Paginator(object_list, paginate_by)

            try:
                page = int(self.request.GET.get('page'))
            except (TypeError, ValueError):
                page = 1
            try:
                page = paginator.page(page)
            except (EmptyPage, InvalidPage):
                page = paginator.page(paginator.num_pages)

            context.update({
                'object_list': page.object_list,
                'page': page,
                'paginator': paginator,
                })

        return context

    def get(self, request, *args, **kwargs):
        self.object_list = self.get_queryset()

        context = {}

        self.object_list, response = self.handle_search_form(
            context, self.object_list)
        if response:
            return response

        response = self.handle_batch_form(context, self.object_list)
        if response:
            return response

        context.update(self.get_context_data(object_list=self.object_list))
        return self.render_to_response(context)

    def handle_search_form(self, context, queryset=None):
        """
        Must return a tuple consisting of a queryset and either a HttpResponse
        or ``None``
        """

        if queryset is None:
            queryset = self.get_query_set(self.request)

        if self.search_form:
            form = self.search_form(self.request.GET, request=self.request)
            if not form.is_valid():
                messages.error(self.request,
                    _('The search query was invalid.'))

                if self.request.get_full_path().endswith('?clear=1'):
                    # No redirect loop generation
                    raise ImproperlyConfigured(
                        'Search form %r does not validate after'
                        ' clearing.' % form)

                return queryset, HttpResponseRedirect('?clear=1')

            queryset = safe_queryset_and(queryset, form.queryset(self.model))

            context['search_form'] = form

        return queryset, None

    def get_batch_actions(self):
        """
        Returns a list of batch action tuples ``(key, name, handler_fn)``

        * ``key``: Something nice, such as ``delete_selected``.
        * ``name``: Will be shown in the dropdown.
        * ``handler_fn``: Callable. Receives the request and the queryset.
        """
        return [
            ('delete_selected', _('Delete selected'), self.delete_selected),
            ]

    def handle_batch_form(self, ctx, queryset):
        """
        May optionally return a HttpResponse which is directly returned to the
        browser
        """

        actions = self.get_batch_actions()
        if not actions:
            return

        class _Form(BatchForm):
            def __init__(self, *args, **kwargs):
                super(_Form, self).__init__(*args, **kwargs)
                self.actions = actions
                self.fields['action'] = forms.ChoiceField(
                    label=_('Action'),
                    choices=[('', '---------')] + [row[:2] for row in actions],
                    widget=forms.HiddenInput,
                    )

            def process(self):
                action = self.cleaned_data.get('action')
                for row in actions:
                    if action == row[0]:
                        return row[2](self.request, self.batch_queryset)

        form = _Form(self.request, queryset)
        ctx['batch_form'] = form

        if form.should_process():
            result = form.process()

            if isinstance(result, HttpResponse):
                return result

            elif hasattr(result, '__iter__'):
                messages.success(self.request,
                    _('Processed the following items:
\n %s') % (
                        u'
\n '.join(
                            unicode(item) for item in result)))

            elif result is not None:
                # Not None, but cannot make sense of it either.
                raise TypeError(u'Return value %r of %s.process() invalid.' % (
                    result,
                    form.__class__.__name__,
                    ))

            info = (
                self.model._meta.app_label,
                self.model._meta.module_name,
                )
            url = tryreverse('%s_%s_list' % info)
            return HttpResponseRedirect(url if url else '.')

    def batch_action_hidden_fields(self, queryset, additional=[]):
        """
        Returns a blob of HTML suitable for jumping back into the batch
        action handler. Most useful for batch action handlers needing to
        present a confirmation and/or form page to the user.

        See ``delete_selected`` below for the usage.
        """
        post_values = [('batchform', 1)] + additional + [
            ('batch_%s' % item.pk, '1') for item in queryset]

        return u'\n'.join(
            u'<input type="hidden" name="%s" value="%s">' % item
            for item in post_values)

    def delete_selected(self, request, queryset):
        # HACK ALARM. XXX ugly. Makes ModelView.add_message never add a
        # message again in this request/response cycle.
        class Bla(list):
            def __contains__(self, key):
                return True
        request._towel_add_message_ignore = Bla()

        allowed = [
            self.deletion_allowed(request, item)
            for item in queryset]
        queryset = [item for item, perm in zip(queryset, allowed) if perm]

        if not queryset:
            messages.error(request,
                _('You are not allowed to delete any object in the selection.'))
            return

        elif not all(allowed):
            messages.warning(request,
                _('Deletion of some objects not allowed. Those have been'
                    ' excluded from the selection already.'))

        if 'confirm' in request.POST:
            messages.success(request, _('Deletion successful.'))
            # Call all delete() methods individually
            [item.delete() for item in queryset]
            return

        context = super(ListView, self).get_context_data(
            title=_('Delete selected'),
            action_queryset=queryset,
            action_hidden_fields=self.batch_action_hidden_fields(
                queryset, [
                    ('batch-action', 'delete_selected'),
                    ('confirm', 1),
                    ]),
            )
        self.template_name_suffix = '_action'
        return self.render_to_response(context)

class DetailView(ModelResourceView):
    template_name_suffix = '_detail'

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()
        context = self.get_context_data(object=self.object)
        return self.render_to_response(context)

class FormView(ModelResourceView):
    form_class = forms.ModelForm
    object = None
    template_name_suffix = '_form'

    def get_form_kwargs(self, **kwargs):
        kw = {'instance': self.object}
        if self.request.method in ('POST', 'PUT'):
            kw.update({
                'data': self.request.POST,
                'files': self.request.FILES,
                })
        kw.update(kwargs)
        return kw

    def get_form_class(self):
        return modelform_factory(self.model,
            form=self.form_class,
            formfield_callback=towel_formfield_callback,
            )

    def get_form(self):
        return self.get_form_class()(**self.get_form_kwargs())

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()
        form = self.get_form()
        context = self.get_context_data(form=form, object=self.object)
        return self.render_to_response(context)

    def post(self, request, *args, **kwargs):
        self.object = self.get_object()
        form = self.get_form()
        if form.is_valid():
            # TODO messages.success
            self.object = form.save()
            # TODO continue editing?
            return redirect(self.object)

        context = self.get_context_data(form=form, object=self.object)
        return self.render_to_response(context)

class LiveFormView(FormView):
    form_class = forms.ModelForm

    def post(self, request, *args, **kwargs):
        self.object = self.get_object()
        form_class = self.get_form_class()

        data = model_to_dict(self.object,
            fields=form_class._meta.fields,
            exclude=form_class._meta.exclude,
            )
        for key, value in request.POST.items():
            data[key] = value

        form = form_class(**self.get_form_kwargs(data=data))

        if form.is_valid():
            self.object = form.save()

            # TODO make this pluggable
            regions = {}
            context = self.get_context_data(
                object=self.object,
                regions=regions,
                )
            self.render_to_response(context)
            return HttpResponse(
                json.dumps(changed_regions(regions, form.changed_data)),
                content_type='application/json')

        return HttpResponse(unicode(form.errors))

class DeleteView(ModelResourceView):
    template_name_suffix = '_delete_confirmation'
    form_class = forms.Form

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()

        if not self.allow('delete', self.object):
            return redirect(self.object)

        if request.method == 'POST':
            form = self.form_class(request.POST)

            if form.is_valid():
                # TODO messages.success()
                self.object.delete()
                return self.object.urls.url('list')
        else:
            form = self.form_class()

        context = self.get_context_data(
            object=self.object,
            form=form,
            )
        return self.render_to_response(context)

class MultitenancyMixin(object):
    def get_queryset(self):
        if self.queryset is not None:
            return safe_queryset_and(self.queryset,
                self.queryset.model._default_manager.for_access(
                    self.request.access))
        elif self.model is not None:
            return self.model._default_manager.for_access(
                self.request.access)
        else:
            raise ImproperlyConfigured("'%s' must define 'queryset' or 'model'"
                                       % self.__class__.__name__)

    def get_form_kwargs(self, **kwargs):
        kwargs['request'] = self.request
        return super(MultitenancyMixin, self).get_form_kwargs(**kwargs)
<commit_msg>Add an easy EditView and AddView to towel.resources<commit_after>"""
This is an experiment in splitting up the monolithic model view class into
smaller, more reusable parts, and using class-based views at the same time
(not the generic class based views, though)
"""

import json

from django import forms
from django.contrib import messages
from django.core.exceptions import ImproperlyConfigured
from django.forms.models import modelform_factory, model_to_dict
from django.http import Http404, HttpResponse, HttpResponseRedirect
from django.shortcuts import get_object_or_404, redirect
from django.utils.translation import ugettext as _
from django.views.generic.base import TemplateView

from towel.forms import BatchForm, towel_formfield_callback
from towel.paginator import Paginator, EmptyPage, InvalidPage
from towel.utils import changed_regions, safe_queryset_and

class ModelResourceView(TemplateView):
    base_template = 'base.html'
    model = None
    queryset = None
    template_name_suffix = None

    def get_context_data(self, **kwargs):
        opts = self.model._meta
        context = {
            'base_template': self.base_template,
            'verbose_name': opts.verbose_name,
            'verbose_name_plural': opts.verbose_name_plural,
            'view': self,

            # TODO add_url
            # TODO list_url
            # TODO adding_allowed
            }
        context.update(kwargs)
        return context

    def get_template_names(self):
        opts = self.model._meta
        names = [
            '{}/{}{}.html'.format(opts.app_label, opts.module_name,
                self.template_name_suffix),
            'resources/object{}.html'.format(self.template_name_suffix),
            ]
        if self.template_name:
            names.insert(0, self.template_name)
        return names

    def get_queryset(self):
        if self.queryset is not None:
            return self.queryset._clone()
        elif self.model is not None:
            return self.model._default_manager.all()
        else:
            raise ImproperlyConfigured("'%s' must define 'queryset' or 'model'"
                                       % self.__class__.__name__)

    def get_object(self):
        return get_object_or_404(self.get_queryset(), **self.kwargs)

    CREATE = type('CREATE', (), {})
    UPDATE = type('UPDATE', (), {})
    DELETE = type('DELETE', (), {})

    def allow(self, action, object=None, **kwargs):
        if action in (self.CREATE, self.UPDATE):
            return True
        return False

    default_messages = {
        CREATE: (messages.SUCCESS,
            _('The new object has been successfully created.')),
        UPDATE: (messages.SUCCESS,
            _('The object has been successfully updated.')),
        DELETE: (messages.SUCCESS,
            _('The object has been successfully deleted.')),

        'adding_denied': (messages.ERROR,
            _('You are not allowed to add objects.')),
        'editing_denied': (messages.ERROR,
            _('You are not allowed to edit this object.')),
        'deletion_denied': (messages.ERROR,
            _('You are not allowed to delete this object.')),
        'deletion_denied_related': (messages.ERROR,
            _('Deletion not allowed: There are %(pretty_classes)s related to this object.')),
    }

    def add_message(self, message, level=None, **kwargs):
        if message in self.default_messages:
            level, message = self.default_messages[message]
        messages.add_message(self.request, level, message, **kwargs)

class ListView(ModelResourceView):
    paginate_by = None
    search_form = None
    template_name_suffix = '_list'

    def get_paginate_by(self, queryset):
        """
        if self.paginate_all_allowed and self.request.GET.get('all'):
            return None
        """
        return self.paginate_by

    def get_context_data(self, object_list, **kwargs):
        context = super(ListView, self).get_context_data(
            object_list=object_list, **kwargs)

        paginate_by = self.get_paginate_by(object_list)
        if paginate_by:
            paginator = Paginator(object_list, paginate_by)

            try:
                page = int(self.request.GET.get('page'))
            except (TypeError, ValueError):
                page = 1
            try:
                page = paginator.page(page)
            except (EmptyPage, InvalidPage):
                page = paginator.page(paginator.num_pages)

            context.update({
                'object_list': page.object_list,
                'page': page,
                'paginator': paginator,
                })

        return context

    def get(self, request, *args, **kwargs):
        self.object_list = self.get_queryset()

        context = {}

        self.object_list, response = self.handle_search_form(
            context, self.object_list)
        if response:
            return response

        response = self.handle_batch_form(context, self.object_list)
        if response:
            return response

        context.update(self.get_context_data(object_list=self.object_list))
        return self.render_to_response(context)

    def handle_search_form(self, context, queryset=None):
        """
        Must return a tuple consisting of a queryset and either a HttpResponse
        or ``None``
        """

        if queryset is None:
            queryset = self.get_query_set(self.request)

        if self.search_form:
            form = self.search_form(self.request.GET, request=self.request)
            if not form.is_valid():
                messages.error(self.request,
                    _('The search query was invalid.'))

                if self.request.get_full_path().endswith('?clear=1'):
                    # No redirect loop generation
                    raise ImproperlyConfigured(
                        'Search form %r does not validate after'
                        ' clearing.' % form)

                return queryset, HttpResponseRedirect('?clear=1')

            queryset = safe_queryset_and(queryset, form.queryset(self.model))

            context['search_form'] = form

        return queryset, None

    def get_batch_actions(self):
        """
        Returns a list of batch action tuples ``(key, name, handler_fn)``

        * ``key``: Something nice, such as ``delete_selected``.
        * ``name``: Will be shown in the dropdown.
        * ``handler_fn``: Callable. Receives the request and the queryset.
        """
        return [
            ('delete_selected', _('Delete selected'), self.delete_selected),
            ]

    def handle_batch_form(self, ctx, queryset):
        """
        May optionally return a HttpResponse which is directly returned to the
        browser
        """

        actions = self.get_batch_actions()
        if not actions:
            return

        class _Form(BatchForm):
            def __init__(self, *args, **kwargs):
                super(_Form, self).__init__(*args, **kwargs)
                self.actions = actions
                self.fields['action'] = forms.ChoiceField(
                    label=_('Action'),
                    choices=[('', '---------')] + [row[:2] for row in actions],
                    widget=forms.HiddenInput,
                    )

            def process(self):
                action = self.cleaned_data.get('action')
                for row in actions:
                    if action == row[0]:
                        return row[2](self.request, self.batch_queryset)

        form = _Form(self.request, queryset)
        ctx['batch_form'] = form

        if form.should_process():
            result = form.process()

            if isinstance(result, HttpResponse):
                return result

            elif hasattr(result, '__iter__'):
                messages.success(self.request,
                    _('Processed the following items:
\n %s') % (
                        u'
\n '.join(
                            unicode(item) for item in result)))

            elif result is not None:
                # Not None, but cannot make sense of it either.
                raise TypeError(u'Return value %r of %s.process() invalid.' % (
                    result,
                    form.__class__.__name__,
                    ))

            info = (
                self.model._meta.app_label,
                self.model._meta.module_name,
                )
            url = tryreverse('%s_%s_list' % info)
            return HttpResponseRedirect(url if url else '.')

    def batch_action_hidden_fields(self, queryset, additional=[]):
        """
        Returns a blob of HTML suitable for jumping back into the batch
        action handler. Most useful for batch action handlers needing to
        present a confirmation and/or form page to the user.

        See ``delete_selected`` below for the usage.
        """
        post_values = [('batchform', 1)] + additional + [
            ('batch_%s' % item.pk, '1') for item in queryset]

        return u'\n'.join(
            u'<input type="hidden" name="%s" value="%s">' % item
            for item in post_values)

    def delete_selected(self, request, queryset):
        # HACK ALARM. XXX ugly. Makes ModelView.add_message never add a
        # message again in this request/response cycle.
        class Bla(list):
            def __contains__(self, key):
                return True
        request._towel_add_message_ignore = Bla()

        allowed = [
            self.deletion_allowed(request, item)
            for item in queryset]
        queryset = [item for item, perm in zip(queryset, allowed) if perm]

        if not queryset:
            messages.error(request,
                _('You are not allowed to delete any object in the selection.'))
            return

        elif not all(allowed):
            messages.warning(request,
                _('Deletion of some objects not allowed. Those have been'
                    ' excluded from the selection already.'))

        if 'confirm' in request.POST:
            messages.success(request, _('Deletion successful.'))
            # Call all delete() methods individually
            [item.delete() for item in queryset]
            return

        context = super(ListView, self).get_context_data(
            title=_('Delete selected'),
            action_queryset=queryset,
            action_hidden_fields=self.batch_action_hidden_fields(
                queryset, [
                    ('batch-action', 'delete_selected'),
                    ('confirm', 1),
                    ]),
            )
        self.template_name_suffix = '_action'
        return self.render_to_response(context)

class DetailView(ModelResourceView):
    template_name_suffix = '_detail'

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()
        context = self.get_context_data(object=self.object)
        return self.render_to_response(context)

class FormView(ModelResourceView):
    form_class = forms.ModelForm
    object = None
    template_name_suffix = '_form'

    def get_form_kwargs(self, **kwargs):
        kw = {'instance': self.object}
        if self.request.method in ('POST', 'PUT'):
            kw.update({
                'data': self.request.POST,
                'files': self.request.FILES,
                })
        kw.update(kwargs)
        return kw

    def get_form_class(self):
        return modelform_factory(self.model,
            form=self.form_class,
            formfield_callback=towel_formfield_callback,
            )

    def get_form(self):
        return self.get_form_class()(**self.get_form_kwargs())

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()
        form = self.get_form()
        context = self.get_context_data(form=form, object=self.object)
        return self.render_to_response(context)

    def post(self, request, *args, **kwargs):
        self.object = self.get_object()
        form = self.get_form()
        if form.is_valid():
            # TODO messages.success
            self.object = form.save()
            # TODO continue editing?
            return redirect(self.object)

        context = self.get_context_data(form=form, object=self.object)
        return self.render_to_response(context)

class EditView(FormView):
    pass
class AddView(FormView):
    def get_object(self):
        return None

class LiveFormView(FormView):
    form_class = forms.ModelForm

    def post(self, request, *args, **kwargs):
        self.object = self.get_object()
        form_class = self.get_form_class()

        data = model_to_dict(self.object,
            fields=form_class._meta.fields,
            exclude=form_class._meta.exclude,
            )
        for key, value in request.POST.items():
            data[key] = value

        form = form_class(**self.get_form_kwargs(data=data))

        if form.is_valid():
            self.object = form.save()

            # TODO make this pluggable
            regions = {}
            context = self.get_context_data(
                object=self.object,
                regions=regions,
                )
            self.render_to_response(context)
            return HttpResponse(
                json.dumps(changed_regions(regions, form.changed_data)),
                content_type='application/json')

        return HttpResponse(unicode(form.errors))

class DeleteView(ModelResourceView):
    template_name_suffix = '_delete_confirmation'
    form_class = forms.Form

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()

        if not self.allow('delete', self.object):
            return redirect(self.object)

        if request.method == 'POST':
            form = self.form_class(request.POST)

            if form.is_valid():
                # TODO messages.success()
                self.object.delete()
                return self.object.urls.url('list')
        else:
            form = self.form_class()

        context = self.get_context_data(
            object=self.object,
            form=form,
            )
        return self.render_to_response(context)

class MultitenancyMixin(object):
    def get_queryset(self):
        if self.queryset is not None:
            return safe_queryset_and(self.queryset,
                self.queryset.model._default_manager.for_access(
                    self.request.access))
        elif self.model is not None:
            return self.model._default_manager.for_access(
                self.request.access)
        else:
            raise ImproperlyConfigured("'%s' must define 'queryset' or 'model'"
                                       % self.__class__.__name__)

    def get_form_kwargs(self, **kwargs):
        kwargs['request'] = self.request
        return super(MultitenancyMixin, self).get_form_kwargs(**kwargs)
<|endoftext|>


<commit_before>import os
import cv2
import sys
import numpy
import shutil
import signal
from PIL import Image
import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GdkPixbuf, GObject, GLib, Gio

from modules import detect
from modules import imgutils

# Width and height of the frames Face Manager will process
WIDTH  = 320
HEIGHT = 240
# Codes for OpenCV camera settings
CV_CAP_PROP_FRAME_WIDTH  = 3
CV_CAP_PROP_FRAME_HEIGHT = 4

class FaceManager:
    """
    A class to control Face Manager GUI operations.
    """

    def __init__(self):
        """
        FaceManager constructor

        Load GUI template from face_manager.glade into a Gtk.Builder object.
        Intantiate objects for each necessary GUI element.
        """
        self.gladefile = 'resources/gui/faces.glade'
        self.gtk = Gtk.Builder()
        self.gtk.add_from_file(self.gladefile)
        self.gtk.connect_signals(self)

        self.FaceFrame = self.gtk.get_object("FaceFrame")
        self.NewFaceName = self.gtk.get_object("NewFaceName")
        self.NewFaceButton = self.gtk.get_object("NewFaceButton")
        self.DeleteFaceButton = self.gtk.get_object("DeleteFaceButton")
        self.FaceManagerDatabase = self.gtk.get_object("DatabaseGrid")

        self.camera = None
        self.face_recognizer = None
        self.namedict = None
        self.last_images = list()

        self.init_newface_name()
        self.init_newface_button()
        self.init_deleteface_button()
        self.update_face_manager_database()

        self.MainWindow = self.gtk.get_object("MainWindow")
        self.MainWindow.connect("delete-event", self.close_button_pressed)
        self.MainWindow.show_all()

        self.name = ""
        self.face_samples_recorded = 0
        self.face_samples = list()
        self.get_face_samples = False

        self.init_camera()
        GLib.idle_add(self.camera_preview)

    def update_face_manager_database(self):
        """
        """
        facedatabase = list()
        for (dirpath, dirnames, filenames) in os.walk('faces'):
            facedatabase.extend(dirnames)
            break

        facedatabase = sorted(facedatabase)

        for child in self.FaceManagerDatabase.get_children():
            self.FaceManagerDatabase.remove(child)

        for i, name in enumerate(facedatabase):
            button = Gtk.ToggleButton()
            box = Gtk.Box()
            image = Gtk.Image()
            image.set_from_file('faces/' + name + '/0.jpg')
            label = Gtk.Label(name)
            box.add(image)
            box.pack_start(label, False, True, 15)
            button.add(box)
            button.set_name(name)
            self.FaceManagerDatabase.add(button)
        self.FaceManagerDatabase.show_all()
        self.trainfacerec()

    def init_newface_name(self):
        """
        Connect the method changed_newface_name() to NewFaceName.
        """
        self.NewFaceName.connect("changed", self.changed_newface_name)

    def changed_newface_name(self, entry):
        """
        Verify if the input name is valid.
        """
        name = self.NewFaceName.get_text()
        if all(c.isalpha() or c.isspace() for c in name):
            self.name = name
            self.NewFaceButton.set_sensitive(True)
        else:
            self.NewFaceButton.set_sensitive(False)

    def init_newface_button(self):
        """
        Connect the method clicked_newface_button() to NewFaceButton.
        """
        self.NewFaceButton.connect("clicked", self.clicked_newface_button)

    def clicked_newface_button(self, button):
        """
        Start new face acquisition process.
        """
        self.start_face_samples_acquisition()

    def init_deleteface_button(self):
        """
        Connect the method clicked_deleteface_button() to DeleteFaceButton.
        """
        self.DeleteFaceButton.connect("clicked", self.clicked_deleteface_button)

    def clicked_deleteface_button(self, button):
        """
        Delete all selected faces from dataset.
        """
        self.delete_face_samples()

    def clicked_deleteface_button(self, button):
        """
        Start new face acquisition process.
        """
        self.delete_face_samples()

    def delete_face_samples(self):
        """
        """
        for child in self.FaceManagerDatabase.get_children():
            if child.get_active():
                shutil.rmtree('faces/' + child.get_name())
                self.FaceManagerDatabase.remove(child)

    def start_face_samples_acquisition(self):
        """
        """
        self.face_samples_recorded = 0
        self.NewFaceName.set_sensitive(False)
        self.NewFaceButton.set_sensitive(False)
        self.get_face_samples = True

    def init_camera(self):
        """
        Initiate video capture using the first webcam found.
        Set camera width and height settings.
        """
        self.camera = cv2.VideoCapture(-1)
        self.camera.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH)
        self.camera.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT)

    def camera_preview(self):
        """
        Get a new frame from camera.
        Process this frame.
        Write to an image file, which is awful, but the only way it worked.
        Display the file on the Frame element of the main window.
        """
        retval, frame = self.camera.read()
        frame = imgutils.crop(frame, 70, 320 - 70, 0, 240)

        frame, found, faces = detect.single_cascade(frame, cascade=detect.CASCADE_FACE,
                                                           return_faces=True,
                                                           drawboxes=True)

        if found:

            x, y, w, h = faces[0]
            face_gray = cv2.cvtColor(frame[y:h, x:w], cv2.COLOR_BGR2GRAY)
            self.last_images.append(imgutils.resize(face_gray, 50, 50))

            if self.face_recognizer != None and self.namedict != None and len(self.last_images) > 8:
                detections = list()
                for img in self.last_images[-8:]:
                    label, conf = self.face_recognizer.predict(img)
                    detections.append(label)
                counts = numpy.bincount(detections)
                detectedlabel = numpy.argmax(counts)
                name = self.namedict[detectedlabel] if counts[detectedlabel] > 6 else "Unknown"
                cv2.putText(frame, name, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # cv2.imwrite(".frame.jpg", frame)
        # pixbuf_frame = GdkPixbuf.Pixbuf.new_from_file(".frame.jpg")
        # self.FaceFrame.set_from_pixbuf(pixbuf_frame)

        h, w, d = frame.shape
        frame_show = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pixbuf = GdkPixbuf.Pixbuf.new_from_data(frame_show.flatten(), GdkPixbuf.Colorspace.RGB, False, 8, w, h, w*3, None, None)
        self.FaceFrame.set_from_pixbuf(pixbuf)

        if self.get_face_samples:

            if found:
                self.face_samples_recorded += 1
                x, y, w, h = faces[0]
                self.face_samples.append(frame[y:h, x:w])

            if self.face_samples_recorded+1 > 50:
                self.save_face_samples()

        return True

    def save_face_samples(self):
        """
        Save captured face samples to our face database
        """
        outputpath = 'faces/' + self.name
        if not os.path.exists(outputpath):
            os.makedirs(outputpath)
        for i, face in enumerate(self.face_samples):
            cv2.imwrite(outputpath + '/' + str(i+1) + '.jpg', imgutils.resize(face, 50, 50))
        cv2.imwrite(outputpath + '/' + str(0) + '.jpg', imgutils.resize(self.face_samples[0], 50, 50))
        self.clean_after_acquisition()

    def clean_after_acquisition(self):
        """
        Make cleaning after adding a new face
        """
        self.get_face_samples = False
        self.face_samples_recorded = 0
        self.face_samples = list()
        self.NewFaceName.set_text("")
        self.NewFaceName.set_sensitive(True)
        self.NewFaceButton.set_sensitive(True)
        self.update_face_manager_database()

    def trainfacerec(self):
        facedatabase = list()
        for (dirpath, dirnames, filenames) in os.walk('faces'):
            facedatabase.extend(dirnames)
            break

        facedatabase = sorted(facedatabase)

        self.namedict = dict()
        faces, labels = list(), list()
        for i, name in enumerate(facedatabase):
            images = list()
            for (dirpath, dirnames, filenames) in os.walk('faces/'+name):
                images.extend(filenames)
                break
            for image in images:
                faces.append(cv2.cvtColor(cv2.imread('faces/'+name+'/'+image), cv2.COLOR_BGR2GRAY))
                labels.append(i)
            self.namedict[i] = name

        if len(facedatabase) > 1:
            self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
            self.face_recognizer.train(faces, numpy.array(labels))

    def close_button_pressed(self, widget, event):
        """
        If close button is pressed, better clean our resources.
        Turn off the camera.
        Execute clean() method, for other resources deallocations.
        Then, actually shut down the poor turret with Gtk.main_quit().
        """
        if self.camera:
            self.camera.release()
        clean()
        Gtk.main_quit()

def clean():
    """
    Use this to close the turret's modules when shutting down.
    """
    pass

def sigint_handler(signum, instant):
    """
    Capture SIGINT signal and quit safely.
    """
    clean()
    sys.exit()

if __name__ == "__main__":

    # Activate capture of SIGINT (Ctrl-C)
    signal.signal(signal.SIGINT, sigint_handler)

    try:
        GObject.threads_init()
        f = FaceManager()
        Gtk.main()
    except KeyboardInterrupt:
        pass
<commit_msg>Deleted duplicate method<commit_after>import os
import cv2
import sys
import numpy
import shutil
import signal
from PIL import Image
import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GdkPixbuf, GObject, GLib, Gio

from modules import detect
from modules import imgutils

# Width and height of the frames Face Manager will process
WIDTH  = 320
HEIGHT = 240
# Codes for OpenCV camera settings
CV_CAP_PROP_FRAME_WIDTH  = 3
CV_CAP_PROP_FRAME_HEIGHT = 4

class FaceManager:
    """
    A class to control Face Manager GUI operations.
    """

    def __init__(self):
        """
        FaceManager constructor

        Load GUI template from face_manager.glade into a Gtk.Builder object.
        Intantiate objects for each necessary GUI element.
        """
        self.gladefile = 'resources/gui/faces.glade'
        self.gtk = Gtk.Builder()
        self.gtk.add_from_file(self.gladefile)
        self.gtk.connect_signals(self)

        self.FaceFrame = self.gtk.get_object("FaceFrame")
        self.NewFaceName = self.gtk.get_object("NewFaceName")
        self.NewFaceButton = self.gtk.get_object("NewFaceButton")
        self.DeleteFaceButton = self.gtk.get_object("DeleteFaceButton")
        self.FaceManagerDatabase = self.gtk.get_object("DatabaseGrid")

        self.camera = None
        self.face_recognizer = None
        self.namedict = None
        self.last_images = list()

        self.init_newface_name()
        self.init_newface_button()
        self.init_deleteface_button()
        self.update_face_manager_database()

        self.MainWindow = self.gtk.get_object("MainWindow")
        self.MainWindow.connect("delete-event", self.close_button_pressed)
        self.MainWindow.show_all()

        self.name = ""
        self.face_samples_recorded = 0
        self.face_samples = list()
        self.get_face_samples = False

        self.init_camera()
        GLib.idle_add(self.camera_preview)

    def update_face_manager_database(self):
        """
        """
        facedatabase = list()
        for (dirpath, dirnames, filenames) in os.walk('faces'):
            facedatabase.extend(dirnames)
            break

        facedatabase = sorted(facedatabase)

        for child in self.FaceManagerDatabase.get_children():
            self.FaceManagerDatabase.remove(child)

        for i, name in enumerate(facedatabase):
            button = Gtk.ToggleButton()
            box = Gtk.Box()
            image = Gtk.Image()
            image.set_from_file('faces/' + name + '/0.jpg')
            label = Gtk.Label(name)
            box.add(image)
            box.pack_start(label, False, True, 15)
            button.add(box)
            button.set_name(name)
            self.FaceManagerDatabase.add(button)
        self.FaceManagerDatabase.show_all()
        self.trainfacerec()

    def init_newface_name(self):
        """
        Connect the method changed_newface_name() to NewFaceName.
        """
        self.NewFaceName.connect("changed", self.changed_newface_name)

    def changed_newface_name(self, entry):
        """
        Verify if the input name is valid.
        """
        name = self.NewFaceName.get_text()
        if all(c.isalpha() or c.isspace() for c in name):
            self.name = name
            self.NewFaceButton.set_sensitive(True)
        else:
            self.NewFaceButton.set_sensitive(False)

    def init_newface_button(self):
        """
        Connect the method clicked_newface_button() to NewFaceButton.
        """
        self.NewFaceButton.connect("clicked", self.clicked_newface_button)

    def clicked_newface_button(self, button):
        """
        Start new face acquisition process.
        """
        self.start_face_samples_acquisition()

    def init_deleteface_button(self):
        """
        Connect the method clicked_deleteface_button() to DeleteFaceButton.
        """
        self.DeleteFaceButton.connect("clicked", self.clicked_deleteface_button)

    def clicked_deleteface_button(self, button):
        """
        Start new face acquisition process.
        """
        self.delete_face_samples()

    def delete_face_samples(self):
        """
        """
        for child in self.FaceManagerDatabase.get_children():
            if child.get_active():
                shutil.rmtree('faces/' + child.get_name())
                self.FaceManagerDatabase.remove(child)

    def start_face_samples_acquisition(self):
        """
        """
        self.face_samples_recorded = 0
        self.NewFaceName.set_sensitive(False)
        self.NewFaceButton.set_sensitive(False)
        self.get_face_samples = True

    def init_camera(self):
        """
        Initiate video capture using the first webcam found.
        Set camera width and height settings.
        """
        self.camera = cv2.VideoCapture(-1)
        self.camera.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH)
        self.camera.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT)

    def camera_preview(self):
        """
        Get a new frame from camera.
        Process this frame.
        Write to an image file, which is awful, but the only way it worked.
        Display the file on the Frame element of the main window.
        """
        retval, frame = self.camera.read()
        frame = imgutils.crop(frame, 70, 320 - 70, 0, 240)

        frame, found, faces = detect.single_cascade(frame, cascade=detect.CASCADE_FACE,
                                                           return_faces=True,
                                                           drawboxes=True)

        if found:

            x, y, w, h = faces[0]
            face_gray = cv2.cvtColor(frame[y:h, x:w], cv2.COLOR_BGR2GRAY)
            self.last_images.append(imgutils.resize(face_gray, 50, 50))

            if self.face_recognizer != None and self.namedict != None and len(self.last_images) > 8:
                detections = list()
                for img in self.last_images[-8:]:
                    label, conf = self.face_recognizer.predict(img)
                    detections.append(label)
                counts = numpy.bincount(detections)
                detectedlabel = numpy.argmax(counts)
                name = self.namedict[detectedlabel] if counts[detectedlabel] > 6 else "Unknown"
                cv2.putText(frame, name, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # cv2.imwrite(".frame.jpg", frame)
        # pixbuf_frame = GdkPixbuf.Pixbuf.new_from_file(".frame.jpg")
        # self.FaceFrame.set_from_pixbuf(pixbuf_frame)

        h, w, d = frame.shape
        frame_show = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pixbuf = GdkPixbuf.Pixbuf.new_from_data(frame_show.flatten(), GdkPixbuf.Colorspace.RGB, False, 8, w, h, w*3, None, None)
        self.FaceFrame.set_from_pixbuf(pixbuf)

        if self.get_face_samples:

            if found:
                self.face_samples_recorded += 1
                x, y, w, h = faces[0]
                self.face_samples.append(frame[y:h, x:w])

            if self.face_samples_recorded+1 > 50:
                self.save_face_samples()

        return True

    def save_face_samples(self):
        """
        Save captured face samples to our face database
        """
        outputpath = 'faces/' + self.name
        if not os.path.exists(outputpath):
            os.makedirs(outputpath)
        for i, face in enumerate(self.face_samples):
            cv2.imwrite(outputpath + '/' + str(i+1) + '.jpg', imgutils.resize(face, 50, 50))
        cv2.imwrite(outputpath + '/' + str(0) + '.jpg', imgutils.resize(self.face_samples[0], 50, 50))
        self.clean_after_acquisition()

    def clean_after_acquisition(self):
        """
        Make cleaning after adding a new face
        """
        self.get_face_samples = False
        self.face_samples_recorded = 0
        self.face_samples = list()
        self.NewFaceName.set_text("")
        self.NewFaceName.set_sensitive(True)
        self.NewFaceButton.set_sensitive(True)
        self.update_face_manager_database()

    def trainfacerec(self):
        facedatabase = list()
        for (dirpath, dirnames, filenames) in os.walk('faces'):
            facedatabase.extend(dirnames)
            break

        facedatabase = sorted(facedatabase)

        self.namedict = dict()
        faces, labels = list(), list()
        for i, name in enumerate(facedatabase):
            images = list()
            for (dirpath, dirnames, filenames) in os.walk('faces/'+name):
                images.extend(filenames)
                break
            for image in images:
                faces.append(cv2.cvtColor(cv2.imread('faces/'+name+'/'+image), cv2.COLOR_BGR2GRAY))
                labels.append(i)
            self.namedict[i] = name

        if len(facedatabase) > 1:
            self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
            self.face_recognizer.train(faces, numpy.array(labels))

    def close_button_pressed(self, widget, event):
        """
        If close button is pressed, better clean our resources.
        Turn off the camera.
        Execute clean() method, for other resources deallocations.
        Then, actually shut down the poor turret with Gtk.main_quit().
        """
        if self.camera:
            self.camera.release()
        clean()
        Gtk.main_quit()

def clean():
    """
    Use this to close the turret's modules when shutting down.
    """
    pass

def sigint_handler(signum, instant):
    """
    Capture SIGINT signal and quit safely.
    """
    clean()
    sys.exit()

if __name__ == "__main__":

    # Activate capture of SIGINT (Ctrl-C)
    signal.signal(signal.SIGINT, sigint_handler)

    try:
        GObject.threads_init()
        f = FaceManager()
        Gtk.main()
    except KeyboardInterrupt:
        pass
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
""" OAuth support, facebook and google.

Session variables:

user.biv_id: User's id, survives after log out
user.oauth_type, User's oauth_type, cleared during log out
user.is_logged_in, User's logged in status
user.display_name, User's full name

oauth.next_uri, redirect uri after successful login,
    cleared after authorize_complete()
oauth.state, oauth state var, cleared after authorize_complete()
oauth.<oauth_type>.token, user's oauth token, cleared during log out

    :copyright: Copyright (c) 2014 Bivio Software, Inc.  All Rights Reserved.
    :license: Apache, see LICENSE for more details.
"""

import flask
import flask_oauthlib.client
from publicprize import controller
from publicprize.auth.model import User
import werkzeug

def _oauth_provider(name, **kwargs):
    """Creates a OAuth client for the named provider."""
    config = controller.app().config
    return flask_oauthlib.client.OAuth(controller.app()).remote_app(
        name,
        consumer_key=config['PP_{}_APP_ID'.format(name.upper())],
        consumer_secret=config['PP_{}_APP_SECRET'.format(name.upper())],
        request_token_url=None,
        **kwargs
    )

_OAUTH_PROVIDER = {
    'facebook': _oauth_provider(
        'facebook',
        request_token_params={'scope': 'email'},
        base_url='https://graph.facebook.com',
        access_token_url='/oauth/access_token',
        authorize_url='https://www.facebook.com/dialog/oauth'
    ),
    'google': _oauth_provider(
        'google',
        request_token_params={
            'scope': 'https://www.googleapis.com/auth/userinfo.email'
        },
        base_url='https://www.googleapis.com/oauth2/v1/',
        access_token_method='POST',
        access_token_url='https://accounts.google.com/o/oauth2/token',
        authorize_url='https://accounts.google.com/o/oauth2/auth',
    )
}

_OAUTH_PROVIDER_DATA_PATH = {
    'facebook': 'me',
    'google': 'userinfo',
}

def add_user_to_session(user):
    """Store user info on session"""
    controller.db.session.add(user)
    controller.db.session.flush()
    flask.session['user.biv_id'] = user.biv_id
    flask.session['user.oauth_type'] = user.oauth_type
    flask.session['user.is_logged_in'] = True
    flask.session['user.display_name'] = user.display_name

def authorize(oauth_type, callback):
    """Return to the "next" or referrer page when return from callback"""
    flask.session['oauth.next_uri'] = flask.request.args.get('next') \
        or flask.request.referrer or None
    state = werkzeug.security.gen_salt(64)
    flask.session['oauth.state'] = state
    return _OAUTH_PROVIDER[oauth_type].authorize(
        callback=callback,
        state=state
    )

def authorize_complete(oauth_type):
    """Handle oauth callback, validate the response and update the user
    model."""
    resp = _OAUTH_PROVIDER[oauth_type].authorized_response()
    next_uri = None
    if 'oauth.next_uri' in flask.session:
        next_uri = flask.session['oauth.next_uri']
        del flask.session['oauth.next_uri']
    if not _validate_auth(resp, oauth_type):
        return flask.redirect(next_uri or '/')
    flask.session['oauth.{}.token'.format(oauth_type)] = (
        resp['access_token'], '')
    _user_from_info(
        oauth_type,
        _OAUTH_PROVIDER[oauth_type].get(
            _OAUTH_PROVIDER_DATA_PATH[oauth_type]).data
    )
    return flask.redirect(next_uri or '/')

def logout():
    """Clear the login state from the session and redirect to root."""
    _clear_session()
    flask.flash('You have successfully logged out.')
    return flask.redirect('/')

def _clear_session():
    """Clear the login state from the session"""
    flask.session['user.is_logged_in'] = False
    oauth_type = flask.session.get('user.oauth_type')
    if oauth_type:
        key = 'oauth.{}.token'.format(oauth_type)
        if flask.session.get(key):
            del flask.session[key]
        del flask.session['user.oauth_type']

def _user_from_info(oauth_type, info):
    """Saves oauth provider user info to user model.
    info arg contains email, id, name."""
    controller.app().logger.info(info)
    if not info.get('email'):
        _clear_session()
        flask.flash('Your email must be provided to this App to login.')
        return
    user = User.query.filter_by(
        oauth_type=oauth_type,
        oauth_id=info['id']
    ).first()
    if user:
        user.display_name = info['name']
        user.user_email = info['email']
    else:
        user = User(
            display_name=info['name'],
            user_email=info['email'],
            oauth_type=oauth_type,
            oauth_id=info['id']
        )
    # TODO(pjm): download user avatar
    # facebook: https://graph.facebook.com/{id}/picture?type=square
    # google: image url provided in info.picture field
    add_user_to_session(user)

def _validate_auth(resp, oauth_type):
    """Validates aouth providers's auth response"""
    app = controller.app()

    if resp is None:
        flask.flash('{} has denied access to this App.'.format(oauth_type))
        app.logger.warn(
            'Access denied: {}'.format(flask.request.args))
        return False
    if isinstance(resp, flask_oauthlib.client.OAuthException):
        flask.flash('{} has denied access to this App.'.format(oauth_type))
        app.logger.warn(resp)
        return False
    state = flask.session['oauth.state']
    del flask.session['oauth.state']
    if state != flask.request.args.get('state'):
        flask.flash('{} has denied access to this App.'.format(oauth_type))
        app.logger.warn(
            'Invalid oauth state, expected: {} response: {}'.format(
                state,
                flask.request.args.get('state'))
        )
        return False
    return True

_FACEBOOK = _OAUTH_PROVIDER['facebook']
_GOOGLE = _OAUTH_PROVIDER['google']

@_FACEBOOK.tokengetter
def _get_facebook_oauth_token():
    """Callback for facebook auth"""
    return flask.session.get('oauth.facebook.token')

@_GOOGLE.tokengetter
def _get_google_access_token():
    """Callback for google auth"""
    return flask.session.get('oauth.google.token')
<commit_msg>refactor error messages into _client_error()<commit_after># -*- coding: utf-8 -*-
""" OAuth support, facebook and google.

Session variables:

user.biv_id: User's id, survives after log out
user.oauth_type, User's oauth_type, cleared during log out
user.is_logged_in, User's logged in status
user.display_name, User's full name

oauth.next_uri, redirect uri after successful login,
    cleared after authorize_complete()
oauth.state, oauth state var, cleared after authorize_complete()
oauth.<oauth_type>.token, user's oauth token, cleared during log out

    :copyright: Copyright (c) 2014 Bivio Software, Inc.  All Rights Reserved.
    :license: Apache, see LICENSE for more details.
"""

import flask
import flask_oauthlib.client
from publicprize import controller
from publicprize.auth.model import User
import werkzeug

def _oauth_provider(name, **kwargs):
    """Creates a OAuth client for the named provider."""
    config = controller.app().config
    return flask_oauthlib.client.OAuth(controller.app()).remote_app(
        name,
        consumer_key=config['PP_{}_APP_ID'.format(name.upper())],
        consumer_secret=config['PP_{}_APP_SECRET'.format(name.upper())],
        request_token_url=None,
        **kwargs
    )

_OAUTH_PROVIDER = {
    'facebook': _oauth_provider(
        'facebook',
        request_token_params={'scope': 'email'},
        base_url='https://graph.facebook.com',
        access_token_url='/oauth/access_token',
        authorize_url='https://www.facebook.com/dialog/oauth'
    ),
    'google': _oauth_provider(
        'google',
        request_token_params={
            'scope': 'https://www.googleapis.com/auth/userinfo.email'
        },
        base_url='https://www.googleapis.com/oauth2/v1/',
        access_token_method='POST',
        access_token_url='https://accounts.google.com/o/oauth2/token',
        authorize_url='https://accounts.google.com/o/oauth2/auth',
    )
}

_OAUTH_PROVIDER_DATA_PATH = {
    'facebook': 'me',
    'google': 'userinfo',
}

def add_user_to_session(user):
    """Store user info on session"""
    controller.db.session.add(user)
    controller.db.session.flush()
    flask.session['user.biv_id'] = user.biv_id
    flask.session['user.oauth_type'] = user.oauth_type
    flask.session['user.is_logged_in'] = True
    flask.session['user.display_name'] = user.display_name

def authorize(oauth_type, callback):
    """Call the OAuth provider's server which returns to the callback
    on completion or failer.
    Store the "next" or referrer page to return to after callback."""
    flask.session['oauth.next_uri'] = flask.request.args.get('next') \
        or flask.request.referrer or None
    # state variable is uri randomizer, verfied in authorize_complete()
    state = werkzeug.security.gen_salt(64)
    flask.session['oauth.state'] = state
    return _OAUTH_PROVIDER[oauth_type].authorize(
        callback=callback,
        state=state
    )

def authorize_complete(oauth_type):
    """Handle oauth callback, validate the response and update the user
    model."""
    resp = _OAUTH_PROVIDER[oauth_type].authorized_response()
    next_uri = None
    if 'oauth.next_uri' in flask.session:
        next_uri = flask.session['oauth.next_uri']
        del flask.session['oauth.next_uri']
    if _validate_auth(resp, oauth_type):
        flask.session['oauth.{}.token'.format(oauth_type)] = (
            resp['access_token'], '')
        _user_from_info(
            oauth_type,
            _OAUTH_PROVIDER[oauth_type].get(
                _OAUTH_PROVIDER_DATA_PATH[oauth_type]).data
        )
    return flask.redirect(next_uri or '/')

def logout():
    """Clear the login state from the session and redirect to root."""
    _clear_session()
    flask.flash('You have successfully logged out.')
    return flask.redirect('/')

def _clear_session():
    """Clear the login state from the session"""
    flask.session['user.is_logged_in'] = False
    if 'oauth.state' in flask.session:
        del flask.session['oauth.state']
    oauth_type = flask.session.get('user.oauth_type')
    if oauth_type:
        key = 'oauth.{}.token'.format(oauth_type)
        if flask.session.get(key):
            del flask.session[key]
        del flask.session['user.oauth_type']

def _client_error(oauth_type, message=None):
    _clear_session()
    flask.flash(message \
        or '{} has denied access to this App.'.format(oauth_type))

def _user_from_info(oauth_type, info):
    """Saves oauth provider user info to user model.
    info arg contains email, id, name."""
    controller.app().logger.info(info)
    if not info.get('email'):
        _client_error(
            oauth, 'Your email must be provided to this App to login.')
        return
    user = User.query.filter_by(
        oauth_type=oauth_type,
        oauth_id=info['id']
    ).first()
    if user:
        user.display_name = info['name']
        user.user_email = info['email']
    else:
        user = User(
            display_name=info['name'],
            user_email=info['email'],
            oauth_type=oauth_type,
            oauth_id=info['id']
        )
    # TODO(pjm): download user avatar
    # facebook: https://graph.facebook.com/{id}/picture?type=square
    # google: image url provided in info.picture field
    add_user_to_session(user)

def _validate_auth(resp, oauth_type):
    """Validates oauth providers's auth response"""
    app = controller.app()

    if resp is None or isinstance(resp, flask_oauthlib.client.OAuthException):
        _client_error(oauth_type)
        app.logger.warn(
            'Access denied: {}, resp: {}'.format(flask.request.args, resp))
        return False
    state = flask.session['oauth.state']
    del flask.session['oauth.state']
    if state != flask.request.args.get('state'):
        _client_error(oauth_type)
        app.logger.warn(
            'Invalid oauth state, expected: {} response: {}'.format(
                state,
                flask.request.args.get('state'))
        )
        return False
    return True

_FACEBOOK = _OAUTH_PROVIDER['facebook']
_GOOGLE = _OAUTH_PROVIDER['google']

@_FACEBOOK.tokengetter
def _get_facebook_oauth_token():
    """Callback for facebook auth"""
    return flask.session.get('oauth.facebook.token')

@_GOOGLE.tokengetter
def _get_google_access_token():
    """Callback for google auth"""
    return flask.session.get('oauth.google.token')
<|endoftext|>


<commit_before>from os.path import abspath, dirname, join

import io
import json
import psycopg2
import re
import requests
import sh
import time

PROJECT_FOLDER=dirname(dirname(dirname(abspath(__file__))))
TEST_FOLDER=join(PROJECT_FOLDER, 'tests')
DOCKER_FOLDER=join(TEST_FOLDER, 'docker')

def wait_for(condition):
    for i in range(120):
        if condition():
            return True
        time.sleep(1)

def pg_is_available():
    def condition():
        try:
            return sql('select 1 + 1;')[0][0] == 2
        except Exception:
            return False
    return condition

def es_is_available():
    url = 'http://localhost:9200'

    def condition():
        try:
            return requests.get(url).json()['tagline'] == 'You Know, for Search'
        except Exception:
            return False
    return condition

def load_sql_file(filename):
    f = join(TEST_FOLDER, 'data', filename)
    with open(f, 'r') as handle:
        sh.psql('--username', 'postgres', 'postgres', _in=handle)

def run_sql_test(filename):
    f = join(TEST_FOLDER, 'test', filename)
    with open(f, 'r') as handle:
        with io.StringIO() as buf:
            sh.psql('--username', 'postgres', '--tuples-only', '--quiet', 'postgres', _in=handle, _out=buf)
            return buf.getvalue()

def load_json_file(filename):
    url = 'http://localhost:9200/_bulk'
    f = join(TEST_FOLDER, 'data', filename)
    headers = {'Content-Type': 'application/x-ndjson'}

    with open(f, 'r') as handle:
        body = handle.read().decode('utf-8')
        requests.post(url, headers=headers, data=body)

def sql(statement):
    with psycopg2.connect(user='postgres', dbname='postgres') as conn:
        with conn.cursor() as cursor:
            cursor.execute(statement)
            return cursor.fetchall()
<commit_msg>Python 3 is unicode by default<commit_after>from os.path import abspath, dirname, join

import io
import json
import psycopg2
import re
import requests
import sh
import time

PROJECT_FOLDER=dirname(dirname(dirname(abspath(__file__))))
TEST_FOLDER=join(PROJECT_FOLDER, 'tests')
DOCKER_FOLDER=join(TEST_FOLDER, 'docker')

def wait_for(condition):
    for i in range(120):
        if condition():
            return True
        time.sleep(1)

def pg_is_available():
    def condition():
        try:
            return sql('select 1 + 1;')[0][0] == 2
        except Exception:
            return False
    return condition

def es_is_available():
    url = 'http://localhost:9200'

    def condition():
        try:
            return requests.get(url).json()['tagline'] == 'You Know, for Search'
        except Exception:
            return False
    return condition

def load_sql_file(filename):
    f = join(TEST_FOLDER, 'data', filename)
    with open(f, 'r') as handle:
        sh.psql('--username', 'postgres', 'postgres', _in=handle)

def run_sql_test(filename):
    f = join(TEST_FOLDER, 'test', filename)
    with open(f, 'r') as handle:
        with io.StringIO() as buf:
            sh.psql('--username', 'postgres', '--tuples-only', '--quiet', 'postgres', _in=handle, _out=buf)
            return buf.getvalue()

def load_json_file(filename):
    url = 'http://localhost:9200/_bulk'
    f = join(TEST_FOLDER, 'data', filename)
    headers = {'Content-Type': 'application/x-ndjson'}

    with open(f, 'r') as handle:
        body = handle.read()
        requests.post(url, headers=headers, data=body)

def sql(statement):
    with psycopg2.connect(user='postgres', dbname='postgres') as conn:
        with conn.cursor() as cursor:
            cursor.execute(statement)
            return cursor.fetchall()
<|endoftext|>