<commit_before>#!/usr/bin/env python

import sys, os, subprocess, shutil
import SpImport

XmlIO = SpImport.SpComp2("PyXmlIO", 1)

import PyOpenColorIO as OCIO

if len(sys.argv) < 3:
    print """
    Usage: spi_to_ocio <color_v7_config.xml> <output_dir>

    Convert a proprietary SPI color configuration format (v7 xml) to an OCIO config.

    env PYTHONPATH=/net/homedirs/jeremys/git/ocio/build/ ./src/testbed/spi_to_ocio.py /net/soft_scratch/users/color/colorspace/OCIO/vfx/v1/vfx.OCIO.xml /net/homedirs/jeremys/git/ocio-configs/spi-vfx

    """
    sys.exit(1)

print ""
print "PyOCIO:", OCIO.__file__
print ""

def GetFileMD5(fname):
    if not os.path.exists(fname):
        raise TypeError("File %s does not exist.", fname)
    process = subprocess.Popen(['md5sum',fname], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    returncode = process.wait()
    if returncode:
        raise TypeError("Could not get cksum on file %s" % fname)
    return process.stdout.read().split()[0]

################################################################################

def BuildConfigFromXMLElement(element, lutDir):
    assert(element.getElementType() == 'color_config')

    config = OCIO.Config()

    # Convert roles
    for oldRole, newRole in [ ("linear", OCIO.ROLE_SCENE_LINEAR),
                              ("colortiming", OCIO.ROLE_COLOR_TIMING),
                              ("spm", "spi.spm"),
                              ("picking", OCIO.ROLE_COLOR_PICKING),
                              ("data", OCIO.ROLE_DATA),
                              ("log", OCIO.ROLE_COMPOSITING_LOG) ]:
        if element.hasAttr(oldRole):
            config.setColorSpaceForRole(newRole, element.getAttr(oldRole))

    for i in xrange(element.getNumberOfChildren()):
        child = element.getChild(i)
        if child.getElementType() == 'colorspace':
            colorspace = BuildColorspaceFromXMLElement(child, lutDir)
            config.addColorSpace(colorspace)
        elif child.getElementType() == 'desc':
            description = child.getText(0)
            config.setDescription(description)
        else:
            print 'TODO: handle config child',child.getElementType()

    return config

def BuildColorTransform(elementArray, lutDir):
    group = OCIO.GroupTransform()

    for element in elementArray:
        attrDict = element.getAttrDict()

        if element.getElementType() == 'null':
            continue

        elif element.getElementType() in ('lut', 'colormatrix'):
            transform = OCIO.FileTransform()

            fname = attrDict.pop('file')
            if not os.path.exists(fname):
                raise Exception("File %s does not exist." % (fname))

            # TODO: COPY AND RENAME FILE!
            base, extension = os.path.basename(fname).rsplit('.',1)
            if extension == 'lut':
                extension = 'spi1d'
            elif extension == 'lut3d':
                extension = 'spi3d'
            elif extension == 'dat':
                extension = 'spimtx'

            newname = os.path.join(lutDir, '%s.%s' % (base, extension))
            if os.path.exists( newname ):
                cksum1 = GetFileMD5(newname)
                cksum2 = GetFileMD5(fname)
                if cksum1 != cksum2:
                    raise Exception("Duplicate files with different contents %s %s." % (newname, fname))
                print '    matched %s %s' % (cksum1, fname)
            else:
                shutil.copy(fname, newname)
                print '    copied %s -> %s' % (fname,newname)

            transform.setSrc('%s.%s' % (base, extension))

            direction = attrDict.pop('direction')
            if direction == 'inverse':
                transform.setDirection(OCIO.TRANSFORM_DIR_INVERSE)
            else:
                transform.setDirection(OCIO.TRANSFORM_DIR_FORWARD)

            interp = attrDict.pop('interpolation', None)
            if interp is None:
                if extension in ('spi3d','3dl'):
                    transform.setInterpolation(OCIO.INTERP_LINEAR)
                elif extension in ('spi1d'):
                    transform.setInterpolation(OCIO.INTERP_NEAREST)
                else:
                    pass
            elif interp == 'linear':
                transform.setInterpolation(OCIO.INTERP_LINEAR)
            elif interp == 'nearest':
                transform.setInterpolation(OCIO.INTERP_NEAREST)
            else:
                print "Unkwnown interpolation",interp

            if attrDict:
                print 'TODO: Handle attrs',attrDict, element.getElementType()

            group.push_back(transform)

        else:
            print 'Unknown color transform',element.getElementType()

    return group

def BuildColorspaceFromXMLElement(element, lutDir):
    attrDict = element.getAttrDict()

    cs = OCIO.ColorSpace()
    cs.setName( attrDict.pop('name') )
    cs.setFamily( attrDict.pop('prefix') )

    if attrDict.pop('type', None) == 'data':
        cs.setIsData(True)

    bitdepth = attrDict.pop('bitdepth')
    if bitdepth == '8':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT8)
    elif bitdepth == '10':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT10)
    elif bitdepth == '12':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT12)
    elif bitdepth == '16':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT16)
    elif bitdepth == '0':
        cs.setBitDepth(OCIO.BIT_DEPTH_F32)
    else:
        #cs.setBitDepth(OCIO.BIT_DEPTH_UNKNOWN)
        raise RuntimeError("Unknown bit depth")

    gpuallocation = attrDict.pop('gpuallocation', None)
    if gpuallocation is None:
        pass
    elif gpuallocation == 'log2':
        cs.setGPUAllocation(OCIO.GPU_ALLOCATION_LG2)
    elif gpuallocation == 'uniform':
        cs.setGPUAllocation(OCIO.GPU_ALLOCATION_UNIFORM)
    else:
        #cs.setGPUAllocation(OCIO.GPU_ALLOCATION_UNKNOWN)
        raise RuntimeError("Unknown bit allocation")

    gpumin = attrDict.pop('gpumin', None)
    if gpumin is not None:
        cs.setGPUMin(float(gpumin))
    gpumax = attrDict.pop('gpumax', None)
    if gpumax is not None:
        cs.setGPUMax(float(gpumax))

    if attrDict:
        print 'TODO: Handle colorspace attrs',attrDict, cs.getName()

    element_toref = None
    element_fromref = None
    description = ''

    for i in xrange(element.getNumberOfChildren()):
        child = element.getChild(i)
        if child.getElementType() == 'to_linear':
            element_toref = child
        elif child.getElementType() == 'from_linear':
            element_fromref = child
        elif child.getElementType() == 'desc':
            description = child.getText(0)
        else:
            print 'TODO: handle colorspace child element',child.getElementType()

    toref, fromref = [], []

    if element_toref is not None:
        toref = [element_toref.getChild(i) for i in xrange(element_toref.getNumberOfChildren())]
    if element_fromref is not None:
        fromref = [element_fromref.getChild(i) for i in xrange(element_fromref.getNumberOfChildren())]

    # This assumes transforms, if both directions are defined, are perfect inverses of each other.
    if len(toref) >= len(fromref):
        transform = BuildColorTransform(toref, lutDir)
        cs.setTransform(transform, OCIO.COLORSPACE_DIR_TO_REFERENCE)
    else:
        transform = BuildColorTransform(fromref, lutDir)
        cs.setTransform(transform, OCIO.COLORSPACE_DIR_FROM_REFERENCE)

    cs.setDescription(description)

    return cs

################################################################################

INPUT_CONFIG = sys.argv[1]
OUTPUT_DIR = sys.argv[2]

LUT_DIR = os.path.join(sys.argv[2] + '/luts')

try: os.makedirs(OUTPUT_DIR)
except: pass
try: os.makedirs(LUT_DIR)
except: pass

p = XmlIO.Parser()
element = p.parse(INPUT_CONFIG)
config = BuildConfigFromXMLElement( element, LUT_DIR)
config.setResourcePath('luts')

OUTPUT_OCIO_FILE = os.path.join(OUTPUT_DIR, 'config.ocio')
print ''
print 'Writing',OUTPUT_OCIO_FILE

f = file(OUTPUT_OCIO_FILE,'w')
f.write(config.getXML())
f.close()

BIN_DIR = os.path.dirname(os.path.abspath(sys.argv[0]))
ocio_info_cmd = os.path.join(BIN_DIR,'ocio_info.py')

OUTPUT_README_FILE = os.path.join(OUTPUT_DIR, 'README')
print ''
print 'Writing',OUTPUT_README_FILE

process = subprocess.Popen([ocio_info_cmd,OUTPUT_OCIO_FILE,], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
returncode = process.wait()
if returncode:
    raise TypeError("Could not run %s %s" % (ocio_info_cmd,OUTPUT_OCIO_FILE))

f = file(OUTPUT_README_FILE,'w')
f.write(process.stdout.read())
f.close()
<commit_msg>updated spi_to_ocio tool<commit_after>#!/usr/bin/env python

import sys, os, subprocess, shutil
import SpImport

XmlIO = SpImport.SpComp2("PyXmlIO", 1)

import PyOpenColorIO as OCIO

if len(sys.argv) < 3:
    print """
    Usage: spi_to_ocio <color_v7_config.xml> <output_dir>

    Convert a proprietary SPI color configuration format (v7 xml) to an OCIO config.

    env PYTHONPATH=/net/homedirs/jeremys/git/OpenColorIO.js/build/ ./src/testbed/spi_to_ocio.py /shots/smf/home/lib/lut/colorspaces.xml /net/homedirs/jeremys/git/OpenColorIO-configs/smf/

    """
    sys.exit(1)

print ""
print "PyOCIO:", OCIO.__file__
print ""

def GetFileMD5(fname):
    if not os.path.exists(fname):
        raise TypeError("File %s does not exist.", fname)
    process = subprocess.Popen(['md5sum',fname], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    returncode = process.wait()
    if returncode:
        raise TypeError("Could not get cksum on file %s" % fname)
    return process.stdout.read().split()[0]

################################################################################

def BuildConfigFromXMLElement(element, lutDir):
    assert(element.getElementType() == 'color_config')

    config = OCIO.Config()

    # Convert roles
    for oldRole, newRole in [ ("linear", OCIO.ROLE_SCENE_LINEAR),
                              ("colortiming", OCIO.ROLE_COLOR_TIMING),
                              ("spm", "spi.spm"),
                              ("picking", OCIO.ROLE_COLOR_PICKING),
                              ("data", OCIO.ROLE_DATA),
                              ("log", OCIO.ROLE_COMPOSITING_LOG) ]:
        if element.hasAttr(oldRole):
            config.setColorSpaceForRole(newRole, element.getAttr(oldRole))

    for i in xrange(element.getNumberOfChildren()):
        child = element.getChild(i)
        if child.getElementType() == 'colorspace':
            colorspace = BuildColorspaceFromXMLElement(child, lutDir)
            config.addColorSpace(colorspace)
        elif child.getElementType() == 'desc':
            description = child.getText(0)
            config.setDescription(description)
        else:
            print 'TODO: handle config child',child.getElementType()

    return config

def BuildColorTransform(elementArray, lutDir):
    group = OCIO.GroupTransform()

    for element in elementArray:
        attrDict = element.getAttrDict()

        if element.getElementType() == 'null':
            continue

        elif element.getElementType() in ('lut', 'colormatrix'):
            transform = OCIO.FileTransform()

            fname = attrDict.pop('file')
            if not os.path.exists(fname):
                raise Exception("File %s does not exist." % (fname))

            # TODO: COPY AND RENAME FILE!
            base, extension = os.path.basename(fname).rsplit('.',1)
            if extension == 'lut':
                extension = 'spi1d'
            elif extension == 'lut3d':
                extension = 'spi3d'
            elif extension == 'dat':
                extension = 'spimtx'

            newname = os.path.join(lutDir, '%s.%s' % (base, extension))
            if os.path.exists( newname ):
                cksum1 = GetFileMD5(newname)
                cksum2 = GetFileMD5(fname)
                if cksum1 != cksum2:
                    raise Exception("Duplicate files with different contents %s %s." % (newname, fname))
                print '    matched %s %s' % (cksum1, fname)
            else:
                shutil.copy(fname, newname)
                print '    copied %s -> %s' % (fname,newname)

            transform.setSrc('%s.%s' % (base, extension))

            direction = attrDict.pop('direction')
            if direction == 'inverse':
                transform.setDirection(OCIO.TRANSFORM_DIR_INVERSE)
            else:
                transform.setDirection(OCIO.TRANSFORM_DIR_FORWARD)

            interp = attrDict.pop('interpolation', None)
            if interp is None:
                if extension in ('spi3d','3dl'):
                    transform.setInterpolation(OCIO.INTERP_LINEAR)
                elif extension in ('spi1d'):
                    transform.setInterpolation(OCIO.INTERP_NEAREST)
                else:
                    pass
            elif interp == 'linear':
                transform.setInterpolation(OCIO.INTERP_LINEAR)
            elif interp == 'nearest':
                transform.setInterpolation(OCIO.INTERP_NEAREST)
            else:
                print "Unkwnown interpolation",interp

            if attrDict:
                print 'TODO: Handle attrs',attrDict, element.getElementType()

            group.push_back(transform)

        else:
            print 'Unknown color transform',element.getElementType()

    return group

def BuildColorspaceFromXMLElement(element, lutDir):
    attrDict = element.getAttrDict()

    cs = OCIO.ColorSpace()
    cs.setName( attrDict.pop('name') )
    cs.setFamily( attrDict.pop('prefix') )

    if attrDict.pop('type', None) == 'data':
        cs.setIsData(True)

    bitdepth = attrDict.pop('bitdepth')
    if bitdepth == '8':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT8)
    elif bitdepth == '10':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT10)
    elif bitdepth == '12':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT12)
    elif bitdepth == '16':
        cs.setBitDepth(OCIO.BIT_DEPTH_UINT16)
    elif bitdepth == '0':
        cs.setBitDepth(OCIO.BIT_DEPTH_F32)
    else:
        #cs.setBitDepth(OCIO.BIT_DEPTH_UNKNOWN)
        raise RuntimeError("Unknown bit depth")

    gpuallocation = attrDict.pop('gpuallocation', None)
    if gpuallocation is None:
        pass
    elif gpuallocation == 'log2':
        cs.setGpuAllocation(OCIO.GPU_ALLOCATION_LG2)
    elif gpuallocation == 'uniform':
        cs.setGpuAllocation(OCIO.GPU_ALLOCATION_UNIFORM)
    else:
        raise RuntimeError("Unknown bit allocation")

    gpumin = attrDict.pop('gpumin', None)
    if gpumin is not None:
        cs.setGpuMin(float(gpumin))
    gpumax = attrDict.pop('gpumax', None)
    if gpumax is not None:
        cs.setGpuMax(float(gpumax))

    if attrDict:
        print 'TODO: Handle colorspace attrs',attrDict, cs.getName()

    element_toref = None
    element_fromref = None
    description = ''

    for i in xrange(element.getNumberOfChildren()):
        child = element.getChild(i)
        if child.getElementType() == 'to_linear':
            element_toref = child
        elif child.getElementType() == 'from_linear':
            element_fromref = child
        elif child.getElementType() == 'desc':
            description = child.getText(0)
        else:
            print 'TODO: handle colorspace child element',child.getElementType()

    toref, fromref = [], []

    if element_toref is not None:
        toref = [element_toref.getChild(i) for i in xrange(element_toref.getNumberOfChildren())]
    if element_fromref is not None:
        fromref = [element_fromref.getChild(i) for i in xrange(element_fromref.getNumberOfChildren())]

    # This assumes transforms, if both directions are defined, are perfect inverses of each other.
    if len(toref) >= len(fromref):
        transform = BuildColorTransform(toref, lutDir)
        cs.setTransform(transform, OCIO.COLORSPACE_DIR_TO_REFERENCE)
    else:
        transform = BuildColorTransform(fromref, lutDir)
        cs.setTransform(transform, OCIO.COLORSPACE_DIR_FROM_REFERENCE)

    cs.setDescription(description)

    return cs

################################################################################

INPUT_CONFIG = sys.argv[1]
OUTPUT_DIR = sys.argv[2]

LUT_DIR = os.path.join(sys.argv[2] + '/luts')

try: os.makedirs(OUTPUT_DIR)
except: pass
try: os.makedirs(LUT_DIR)
except: pass

p = XmlIO.Parser()
element = p.parse(INPUT_CONFIG)
config = BuildConfigFromXMLElement( element, LUT_DIR)
config.setResourcePath('luts')

OUTPUT_OCIO_FILE = os.path.join(OUTPUT_DIR, 'config.ocio')
print ''
print 'Writing',OUTPUT_OCIO_FILE

f = file(OUTPUT_OCIO_FILE,'w')
f.write(config.getXML())
f.close()

BIN_DIR = os.path.dirname(os.path.abspath(sys.argv[0]))
ocio_info_cmd = os.path.join(BIN_DIR,'ocio_info.py')

OUTPUT_README_FILE = os.path.join(OUTPUT_DIR, 'README')
print ''
print 'Writing',OUTPUT_README_FILE

process = subprocess.Popen([ocio_info_cmd,OUTPUT_OCIO_FILE,], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
returncode = process.wait()
if returncode:
    raise TypeError("Could not run %s %s" % (ocio_info_cmd,OUTPUT_OCIO_FILE))

f = file(OUTPUT_README_FILE,'w')
f.write(process.stdout.read())
f.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python

from main import PKT_DIR_INCOMING, PKT_DIR_OUTGOING
import socket, struct
from bisect import bisect_left
from datetime import datetime

# TODO: Feel free to import any Python standard modules as necessary.
# (http://docs.python.org/2/library/)
# You must NOT use any 3rd-party libraries, though.

class Firewall:
    def __init__(self, config, timer, iface_int, iface_ext):
        self.timer = timer
        self.iface_int = iface_int
        self.iface_ext = iface_ext

        parser = RulesParser(config['rule'])
        self.rules = parser.parse_rules()

        # TODO: Also do some initialization if needed.

    def handle_timer(self):
        # TODO: For the timer feature, refer to bypass.py
        pass

    # @pkt_dir: either PKT_DIR_INCOMING or PKT_DIR_OUTGOING
    # @pkt: the actual data of the IPv4 packet (including IP header)
    def handle_packet(self, pkt_dir, pkt):
        # TODO: Your main firewall code will be here.

        protocol, ext_IP_address, ext_port, is_dns_pkt, domain_name = self.read_packet(pkt, pkt_dir)
        wrapped_packet = Wrap_Packet(protocol, ext_IP_address, ext_port, is_dns_pkt, domain_name)

        verdict = self.packet_lookup(wrapped_packet)

        if verdict == "pass":
          if pkt_dir == PKT_DIR_INCOMING:
            self.iface_int.send_ip_packet(pkt)
          else: # pkt_dir == PKT_DIR_OUTGOING
            self.iface_ext.send_ip_packet(pkt)

    # Acts as a parser for the packet
    # Returns the protocol, external IP address, and the external port associated with the packet
    # Also determines whether or not a packet is a DNS packet and returns that as well
    def read_packet(self, pkt, pkt_dir):
      # Need to retrieve the protocol that the packet follows

      if pkt_dir == PKT_DIR_INCOMING:
        pass
        # Retrieve the source IP address and source port of the packet
      else: # pkt_dir == PKT_DIR_OUTGOING
        pass
        # Retrieve the destination IP address and destination port of the packet

      # Determine whether or not the packet is a DNS query

      return protocol, ext_IP_address, ext_port, is_dns_pkt, domain_name

    # Looks through the self.rules list and returns the verdict of the latest
    # rule in the list that matches the packet fields
    # Returns verdict==True if no rules in the list match the packet fields
    def packet_lookup(self, wrapped_packet):

    	# Set verdict initially to true so that if no rules match packet fields
    	# then the packet will be passed to the appropriate interface
    	verdict = True

    	for rule in self.rules:
    		if rule.protocol == wrapped_packet.protocol:
    			# Examine rule further since rule protocol matches packet protocol (TCP/UDP/ICMP)
    			if wrapped_packet.ext_IP_address == rule.ext_IP_address and wrapped_packet.ext_port == rule.ext_port:
    				verdict = rule.verdict
    		elif rule.protocol == "dns" and wrapped_packet.is_dns_pkt == true:
    			# Examine rule further since it is a DNS rule and packet is dns query
    			# wrapped_packet only has a domain_name field if is_dns_pkt is true for wrapped_packet
    			if wrapped_packet.domain_name == rule.domain_name:
    				verdict = rule.verdict
    	return verdict

    # TODO: You can add more methods as you want.

class RulesParser:

  def __init__(self, filename):
    self.filename = filename

  def parse_rules(self):
    f = open(self.filename, 'r')
    rules = []
    for line in f:
      line = line.strip()
      #Ignore All lines that are blank or comment lines
      if line and not self._is_comment_line(line):
        rule = self.parse_line(line)
        if rule:
          rules.append(rule)
    return rules

  def parse_line(self, line):
    tokens = line.split()
    #Make all tokens lowercase so we ignore case sensitivity
    tokens = map(lambda token: token.lower(), tokens)

    #Rules that have 4 fields are normal rules
    if len(tokens) == 4:
      return Rule(*tokens)
    #DNS Rules
    elif len(tokens) == 3:
      return DNSRule(*tokens)
    else:
      return None

  def _is_comment_line(self, line):
    return line.startswith("%")

class GeoDBParser:

  def __init__(self, filename):
    self.filename = filename

  def parse_lines(self):
    f = open(self.filename, 'r')
    nodes = []
    for line in f:
      node = self.parse_line(line)
      nodes.append(node)
    return nodes

  def parse_line(self, line):
    tokens = line.split()
    #Make all tokens lowercase so we ignore case sensitivity
    tokens = map(lambda token: token.lower(), tokens)

    #GeoDB lines should have 3 fields
    if len(tokens) == 3:
      return GeoDBNode(*tokens)
    else:
      return None

class GeoDBNode:

  def __init__(self, start_ip, end_ip, country_code):
    self.start_ip = self.ip_to_int(start_ip)
    self.end_ip = self.ip_to_int(end_ip)
    self.country_code = country_code

  def __lt__(self, other):
    ip = self.ip_to_int(other)
    return self.end_ip < ip

  def __le__(self, other):
    return self.__lt__(other) or self.__eq__(other)

  def __ne__(self, other):
    return not self.__eq__(other)

  def __eq__(self, other):
    ip = self.ip_to_int(other)
    return self.start_ip <= ip and ip <= self.end_ip

  def __gt__(self, other):
    ip = self.ip_to_int(other)
    return self.start_ip > ip

  def __ge__(self, other):
    return self.__gt__(other) or self.__eq__(other)

  def ip_to_int(self, ip):
    #From http://gist/githib.com/cslarsen/1595135
    return reduce(lambda a,b: a<<8 | b, map(int, ip.split(".")))

class Rule:

  def __init__(self, verdict, protocol, ext_IP_address, ext_port):
    self.verdict = verdict
    self.protocol = protocol
    self.ext_IP_address = IPAddressField(ext_IP_address)
    self.ext_port = ExtPortField(ext_port)

class WrappedPacket:
	def __init__(self, protocol, ext_IP_address, ext_port, is_dns_pkt, domain_name):
		self.protocol = protocol
    self.ext_IP_address = IPAddressField(ext_IP_address)
    self.ext_port = ExtPortField(ext_port)
    self.is_dns_pkt = is_dns_pkt
    if self.is_dns_pkt:
    	self.domain_name = DomainNameField(domain_name)

class DNSRule(Rule):

  def __init__(self, verdict, protocol, domain_name):
    self.verdict = verdict
    self.protocol = protocol
    self.domain_name = DomainNameField(domain_name)

class IPAddressField:

  geoParser = GeoDBParser('geoipdb.txt')
  geo_nodes = geoParser.parse_lines()

  def __init__(self, ext_IP_address):
  	self.ext_IP_address = ext_IP_address

  	# Have this set up as a field so that way don't have to
  	# keep determining if IPAddressField obj wraps an IP prefix
  	if "/" in self.ext_IP_address:
  		self.is_IP_prefix = true
  		# Element 0 is the IP, Element 1 is the slash #
  		_decimal_ip_and_prefix = other.ext_IP_address.split("/")
  		# Ask: any reason to make ip_to_int an instance function?
  		self._decimal_ip = self.ip_to_int(_decimal_ip_and_prefix[0])
  		self.slash_num = _decimal_ip_and_prefix[1]
  		self.relevant_portion = self.relevant_ip_portion(self._decimal_ip, self.slash_num)
    else:
      self.is_IP_prefix = false

  # Assume that the lhs of "==" is always the external IP address
  # of the packet, while the rhs is always the external Ip addr of the
  # rule that you're trying to match up with the packet
  def __eq__(self, other):
  	if other.ext_IP_address == "any":
  		return True
  	elif len(other.ext_IP_address) == 2:
  		# other.ext_IP_address is a 2-byte country code
  		return belongs_to_country(self.ext_IP_address, other.ext_IP_address)
  	elif other.is_IP_prefix:
  		# TODO: Deal with IP prefix case
  		decimal_ip = ip_to_int(self.ext_IP_address)
  		return self.relevant_ip_portion(decimal_ip, other.slash_num) == other.relevant_portion
  	else:
  		# other.ext_IP_address is just an IP address
  		return self.ext_IP_address == other.ext_IP_address

  #Returns True if the given ip belongs to a certain country. False o/w
  def belongs_to_country(self, ip, country_code):
    geo_node = self._bin_search(ip)
    if geo_node:
      return geo_node.country_code == country_code
    else:
      return False

  #Finds the GeoDBNode within the range of the given IP or None if no node of appropriate range is found
  def _bin_search(self, ip):
    i = bisect_left(self.geo_nodes, ip)
    if i != len(self.geo_nodes):
      node = self.geo_nodes[i]
      if node == ip:
        return node
      else:
        return None
    else:
      return None

  def ip_to_int(self, ip):
    #From http://gist/githib.com/cslarsen/1595135
    return reduce(lambda a,b: a<<8 | b, map(int, ip.split(".")))

  # Return the network component of the inputted IP address zero-extended
  # at the end (for later comparison of network component portions)
  def relevant_ip_portion(self, ip, slash_num):
  	subnet_mask = 0xFFFFFFFF >> (32 - slash_num)
  	subnet_mask = subnet_mask << (32 - slash_num)
    relevant_ip = ip & subnet_mask
  	return relevant_ip

class ExtPortField:

  def __init__(self, ext_port):
  	self.ext_port = ext_port
  	self.is_a_range = false
   	if "-" in self.ext_port:
  		self.is_a_range = true
  		self._temp_list = self.ext_port.split("-")
  		self.start_port = int(self._temp_list[0])
  		self.end_port = int(self._temp_list[1])

  # Assume that the lhs of "==" is always the external port
  # of the packet, while the rhs is always the external port of the
  # rule that you're trying to match up with the packet
  def __eq__(self, other):
  	if other.ext_port == "any":
  		return True
  	elif other.is_a_range == true:
      ext_port_as_int = int(self.ext_port)
  		return ext_port_as_int >= other.start_port and ext_port_as_int <= other.end_port
  	else:
  		# other.ext_port should be a single value
  		return self.ext_port == other.ext_port

  def _is_integer(self, ext_port):
    try:
      int(ext_port)
      return True
    except ValueError:
      return False

class DomainNameField:

  def __init__(self, domain_name):
  	self._domain_name = domain_name
  	self._intermediate_list = self._domain_name.split(".")
  	# self.rev_domain_name_list is the result of a string split into
  	# multiple elements by the . delimiter, followed by a reverse oper.
  	# that reverses the elements in the list
  	self.rev_domain_name_list = reversed(self._intermediate_list)

  # Assume that the lhs of "==" is always the domain name
  # of the packet, while the rhs is always the domain name of
  # the rule that you're trying to match up with the packet
  def __eq__(self, other):
  	for i, partURL in enumerate(other.rev_domain_name_list):
  		# Assume that DNS query rules only have good syntax
  		if partURL == "*":
  			# Since the current partURL is "*", do not care what the rest of self.rev_domain_name_list is
  			return True
  		else:
  			# partURL is a portion of a url like "gov" or "fda"
  			if partURL != self.rev_domain_name_list[i]:
  				return False
  	# Should never reach this part
  	raise Exception("WTF WE FUCKED UP")

# TODO: You may want to add more classes/functions as well.
<commit_msg>Finished packet parsing - Need a domain name converter method<commit_after>#!/usr/bin/env python

from main import PKT_DIR_INCOMING, PKT_DIR_OUTGOING
import socket, struct
from bisect import bisect_left
from datetime import datetime

# TODO: Feel free to import any Python standard modules as necessary.
# (http://docs.python.org/2/library/)
# You must NOT use any 3rd-party libraries, though.

class Firewall:
    def __init__(self, config, timer, iface_int, iface_ext):
        self.timer = timer
        self.iface_int = iface_int
        self.iface_ext = iface_ext

        parser = RulesParser(config['rule'])
        self.rules = parser.parse_rules()

        # TODO: Also do some initialization if needed.

    def handle_timer(self):
        # TODO: For the timer feature, refer to bypass.py
        pass

    # @pkt_dir: either PKT_DIR_INCOMING or PKT_DIR_OUTGOING
    # @pkt: the actual data of the IPv4 packet (including IP header)
    def handle_packet(self, pkt_dir, pkt):
        # TODO: Your main firewall code will be here.

        protocol, ext_IP_address, ext_port, check_dns_rules, domain_name, ok_header_len = self.read_packet(pkt, pkt_dir)
        if (!ok_header_len):
          return
        wrapped_packet = WrappedPacket(protocol, ext_IP_address, ext_port, check_dns_rules, domain_name)

        verdict = self.packet_lookup(wrapped_packet)

        if verdict == "pass":
          if pkt_dir == PKT_DIR_INCOMING:
            self.iface_int.send_ip_packet(pkt)
          else: # pkt_dir == PKT_DIR_OUTGOING
            self.iface_ext.send_ip_packet(pkt)

    # Acts as a parser for the packet
    # Returns the protocol, external IP address, and the external port associated with the packet
    # Also determines whether or not a packet is a DNS packet and returns that as well
    # TODO: Check that packets are laid out in memory BIG-ENDIAN
    def read_packet(self, pkt, pkt_dir):

      # Need to retrieve the protocol that the packet follows
      protocol_tmp = struct.unpack('!B',pkt[9])[0] # Protocol number corresponding to TCP/UDP/ICMP-type

      if protocol_tmp == 1:
        protocol = "ICMP"
      elif protocol_tmp == 6:
        protocol = "TCP"
      else: # protocol_tmp == 17
        protocol = "UDP"

      header_len_tmp = struct.unpack('!B',pkt[0])[0]
      header_len = header_len_tmp & 0xFF
      if header_len < 5:
        # Check additional specs to see that packets with header length < 5 should be dropped.
        ok_header_len = False
      else: # header_len >= 5
        tl_index = header_len*4
        ok_header_len = True

      if pkt_dir == PKT_DIR_INCOMING:
        ext_ip_tmp = pkt[12:16] # external IP address is source IP address
        if protocol != "ICMP":
          ext_port_tmp = struct.unpack('!H',pkt[tl_index:tl_index+2])[0]
        # Retrieve the source IP address and source port of the packet
      else: # pkt_dir == PKT_DIR_OUTGOING
        ext_ip_tmp = pkt[16:20] # external IP address is destination IP address
        if protocol != "ICMP":
          ext_port_tmp = struct.unpack('!H',pkt[tl_index+2:tl_index+4])[0]

      # Packet is ICMP-type packet
      if protocol == "ICMP":
        ext_port_tmp = struct.unpack('!B',pkt[tl_index])[0] # Independent of pkt_dir value

      # Retrieve the string representation of the external IP address of a packet
      external_IP_address = socket.inet_ntoa(ext_ip_tmp)

      # Retrieve the string representation of the external port of a packet
      ext_port = (str) ext_port_tmp

      # check_dns_rules determines whether or not a packet should be considered
      # for DNS rule matching. Initially set to False unless can prove otherwise.
      check_dns_rules = False
      domain_name = None
      dst_port = struct.unpack('!H',pkt[tl_index+2:tl_index+4])[0]

      if protocol == "UDP" and dst_port == 53:
        # Check to see if there is exactly one DNS question entry
        # Application layer starts at index = tl_index + 8
        al_index = tl_index + 8
        QDCOUNT = struct.unpack('!H',pkt[al_index+4:al_index+6])[0]
        if QDCOUNT == 1:
          len_byte_index = al_index+12
          length_byte = struct.unpack('!B',pkt[len_byte_index])[0]
          list_of_domain_parts = []
          while (length_byte != 0):
            tuple_of_ascii_char = struct.unpack('!B',pkt[len_byte_index+1:len_byte_index+length_byte+1])
            list_of_domain_parts.append(tuple_of_ascii_char)
            len_byte_index = len_byte_index+length_byte+1
            length_byte = struct.unpack('!B', pkt[len_byte_index])[0]

          # Converts the list list_of_domain_parts to a string domain name
          domain_name = parse_domain_name(list_of_domain_parts)

          # len_byte_index now represents the starting index of QTYPE
          QTYPE = struct.unpack('!H',pkt[len_byte_index:len_byte_index+2])[0]
          QCLASS = struct.unpack('!H',pkt[len_byte_index+2:len_byte_index+4])[0]

          if ((QTYPE == 1 or QTYPE == 28) and QCLASS == 1):
            check_dns_rules = True

      return protocol, ext_IP_address, ext_port, check_dns_rules, domain_name, ok_header_len

    # Plz write this for me Ben
    # Given a list
    # [(0x77, 0x77, 0x77,), (0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65,), (0x63, 0x6f, 0x6d,)]
    # return "www.google.com"
    def parse_domain_name(list_of_domain_parts):
      pass

    # Looks through the self.rules list and returns the verdict of the latest
    # rule in the list that matches the packet fields
    # Returns verdict==True if no rules in the list match the packet fields
    def packet_lookup(self, wrapped_packet):

    	# Set verdict initially to true so that if no rules match packet fields
    	# then the packet will be passed to the appropriate interface
    	verdict = True

    	for rule in self.rules:
    		if rule.protocol == wrapped_packet.protocol:
    			# Examine rule further since rule protocol matches packet protocol (TCP/UDP/ICMP)
    			if wrapped_packet.ext_IP_address == rule.ext_IP_address and wrapped_packet.ext_port == rule.ext_port:
    				verdict = rule.verdict
    		elif rule.protocol == "dns" and wrapped_packet.check_dns_rules == true:
    			# Examine rule further since check_dns_rules indicates that a packet should be checked against DNS rules
    			# wrapped_packet only has a domain_name field if check_dns_rules is true for wrapped_packet
    			if wrapped_packet.domain_name == rule.domain_name:
    				verdict = rule.verdict
    	return verdict

    # TODO: You can add more methods as you want.

class RulesParser:

  def __init__(self, filename):
    self.filename = filename

  def parse_rules(self):
    f = open(self.filename, 'r')
    rules = []
    for line in f:
      line = line.strip()
      #Ignore All lines that are blank or comment lines
      if line and not self._is_comment_line(line):
        rule = self.parse_line(line)
        if rule:
          rules.append(rule)
    return rules

  def parse_line(self, line):
    tokens = line.split()
    #Make all tokens lowercase so we ignore case sensitivity
    tokens = map(lambda token: token.lower(), tokens)

    #Rules that have 4 fields are normal rules
    if len(tokens) == 4:
      return Rule(*tokens)
    #DNS Rules
    elif len(tokens) == 3:
      return DNSRule(*tokens)
    else:
      return None

  def _is_comment_line(self, line):
    return line.startswith("%")

class GeoDBParser:

  def __init__(self, filename):
    self.filename = filename

  def parse_lines(self):
    f = open(self.filename, 'r')
    nodes = []
    for line in f:
      node = self.parse_line(line)
      nodes.append(node)
    return nodes

  def parse_line(self, line):
    tokens = line.split()
    #Make all tokens lowercase so we ignore case sensitivity
    tokens = map(lambda token: token.lower(), tokens)

    #GeoDB lines should have 3 fields
    if len(tokens) == 3:
      return GeoDBNode(*tokens)
    else:
      return None

class GeoDBNode:

  def __init__(self, start_ip, end_ip, country_code):
    self.start_ip = self.ip_to_int(start_ip)
    self.end_ip = self.ip_to_int(end_ip)
    self.country_code = country_code

  def __lt__(self, other):
    ip = self.ip_to_int(other)
    return self.end_ip < ip

  def __le__(self, other):
    return self.__lt__(other) or self.__eq__(other)

  def __ne__(self, other):
    return not self.__eq__(other)

  def __eq__(self, other):
    ip = self.ip_to_int(other)
    return self.start_ip <= ip and ip <= self.end_ip

  def __gt__(self, other):
    ip = self.ip_to_int(other)
    return self.start_ip > ip

  def __ge__(self, other):
    return self.__gt__(other) or self.__eq__(other)

  def ip_to_int(self, ip):
    #From http://gist/githib.com/cslarsen/1595135
    return reduce(lambda a,b: a<<8 | b, map(int, ip.split(".")))

class Rule:

  def __init__(self, verdict, protocol, ext_IP_address, ext_port):
    self.verdict = verdict
    self.protocol = protocol
    self.ext_IP_address = IPAddressField(ext_IP_address)
    self.ext_port = ExtPortField(ext_port)

class WrappedPacket:
	def __init__(self, protocol, ext_IP_address, ext_port, check_dns_rules, domain_name):
		self.protocol = protocol
    self.ext_IP_address = IPAddressField(ext_IP_address)
    self.ext_port = ExtPortField(ext_port)
    self.check_dns_rules = check_dns_rules
    if self.check_dns_rules:
    	self.domain_name = DomainNameField(domain_name)

class DNSRule(Rule):

  def __init__(self, verdict, protocol, domain_name):
    self.verdict = verdict
    self.protocol = protocol
    self.domain_name = DomainNameField(domain_name)

class IPAddressField:

  geoParser = GeoDBParser('geoipdb.txt')
  geo_nodes = geoParser.parse_lines()

  def __init__(self, ext_IP_address):
  	self.ext_IP_address = ext_IP_address

  	# Have this set up as a field so that way don't have to
  	# keep determining if IPAddressField obj wraps an IP prefix
  	if "/" in self.ext_IP_address:
  		self.is_IP_prefix = true
  		# Element 0 is the IP, Element 1 is the slash #
  		_decimal_ip_and_prefix = other.ext_IP_address.split("/")
  		# Ask: any reason to make ip_to_int an instance function?
  		self._decimal_ip = self.ip_to_int(_decimal_ip_and_prefix[0])
  		self.slash_num = _decimal_ip_and_prefix[1]
  		self.relevant_portion = self.relevant_ip_portion(self._decimal_ip, self.slash_num)
    else:
      self.is_IP_prefix = false

  # Assume that the lhs of "==" is always the external IP address
  # of the packet, while the rhs is always the external Ip addr of the
  # rule that you're trying to match up with the packet
  def __eq__(self, other):
  	if other.ext_IP_address == "any":
  		return True
  	elif len(other.ext_IP_address) == 2:
  		# other.ext_IP_address is a 2-byte country code
  		return belongs_to_country(self.ext_IP_address, other.ext_IP_address)
  	elif other.is_IP_prefix:
  		# TODO: Deal with IP prefix case
  		decimal_ip = ip_to_int(self.ext_IP_address)
  		return self.relevant_ip_portion(decimal_ip, other.slash_num) == other.relevant_portion
  	else:
  		# other.ext_IP_address is just an IP address
  		return self.ext_IP_address == other.ext_IP_address

  #Returns True if the given ip belongs to a certain country. False o/w
  def belongs_to_country(self, ip, country_code):
    geo_node = self._bin_search(ip)
    if geo_node:
      return geo_node.country_code == country_code
    else:
      return False

  #Finds the GeoDBNode within the range of the given IP or None if no node of appropriate range is found
  def _bin_search(self, ip):
    i = bisect_left(self.geo_nodes, ip)
    if i != len(self.geo_nodes):
      node = self.geo_nodes[i]
      if node == ip:
        return node
      else:
        return None
    else:
      return None

  def ip_to_int(self, ip):
    #From http://gist/githib.com/cslarsen/1595135
    return reduce(lambda a,b: a<<8 | b, map(int, ip.split(".")))

  # Return the network component of the inputted IP address zero-extended
  # at the end (for later comparison of network component portions)
  def relevant_ip_portion(self, ip, slash_num):
  	subnet_mask = 0xFFFFFFFF >> (32 - slash_num)
  	subnet_mask = subnet_mask << (32 - slash_num)
    relevant_ip = ip & subnet_mask
  	return relevant_ip

class ExtPortField:

  def __init__(self, ext_port):
  	self.ext_port = ext_port
  	self.is_a_range = false
   	if "-" in self.ext_port:
  		self.is_a_range = true
  		self._temp_list = self.ext_port.split("-")
  		self.start_port = int(self._temp_list[0])
  		self.end_port = int(self._temp_list[1])

  # Assume that the lhs of "==" is always the external port
  # of the packet, while the rhs is always the external port of the
  # rule that you're trying to match up with the packet
  def __eq__(self, other):
  	if other.ext_port == "any":
  		return True
  	elif other.is_a_range == true:
      ext_port_as_int = int(self.ext_port)
  		return ext_port_as_int >= other.start_port and ext_port_as_int <= other.end_port
  	else:
  		# other.ext_port should be a single value
  		return self.ext_port == other.ext_port

  def _is_integer(self, ext_port):
    try:
      int(ext_port)
      return True
    except ValueError:
      return False

class DomainNameField:

  def __init__(self, domain_name):
  	self._domain_name = domain_name
  	self._intermediate_list = self._domain_name.split(".")
  	# self.rev_domain_name_list is the result of a string split into
  	# multiple elements by the . delimiter, followed by a reverse oper.
  	# that reverses the elements in the list
  	self.rev_domain_name_list = reversed(self._intermediate_list)

  # Assume that the lhs of "==" is always the domain name
  # of the packet, while the rhs is always the domain name of
  # the rule that you're trying to match up with the packet
  def __eq__(self, other):
  	for i, partURL in enumerate(other.rev_domain_name_list):
  		# Assume that DNS query rules only have good syntax
  		if partURL == "*":
  			# Since the current partURL is "*", do not care what the rest of self.rev_domain_name_list is
  			return True
  		else:
  			# partURL is a portion of a url like "gov" or "fda"
  			if partURL != self.rev_domain_name_list[i]:
  				return False
  	# Should never reach this part
  	raise Exception("WTF WE FUCKED UP")

# TODO: You may want to add more classes/functions as well.
<|endoftext|>


<commit_before>from flask import Blueprint, request
from flask.views import MethodView
from apps.controllers.base import BaseController
from apps.models.company import CompanyModel

companies = Blueprint('companies', __name__)

class CompanyController(MethodView, BaseController):

    def get(self, company_id):
        """
        Returns in JSON format one or many Companies

        :param company_id:
        """
        company = []
        pagination = None

        if company_id is None:
            page = int(request.args.get('page', '1'))
            count = int(request.args.get('count', '5'))

            pagination = CompanyModel.query.paginate(page, count, False)
            company = pagination.items

            if len(company) == 0:
                return self.response(404)

        else:
            company = CompanyModel.query.filter_by(id=company_id).first()
            if company is None:
                return self.response(404)

        return self.response(200, company, pagination)

    def post(self):
        """
        Create a new Company

        :return: Code 201
        """
        return self.response(201)

    def put(self, company_id):
        """
        Edit a Company

        :param company_id:
        :return: Code 200
        """
        return self.response(202)

    def delete(self, company_id):
        """
        Delete a Company

        :param company_id:
        :return:
        """
        return self.response(204)

companies_view = CompanyController.as_view('companies')
companies.add_url_rule('', view_func=companies_view, methods=['POST'])
companies.add_url_rule('', defaults={'company_id': None}, view_func=companies_view, methods=['GET'])
companies.add_url_rule('/<int:company_id>', view_func=companies_view, methods=['GET', 'PUT', 'DELETE'])<commit_msg>Added POST, PUT and DELETE<commit_after>from flask import Blueprint, request
from flask.views import MethodView
from apps.controllers.base import BaseController
from apps.models.company import CompanyModel

companies = Blueprint('companies', __name__)

class CompanyController(MethodView, BaseController):

    def get(self, company_id):
        """
        Returns in JSON format one or many Companies

        :param company_id:
        """
        pagination = None

        if company_id is None:
            page = int(request.args.get('page', '1'))
            count = int(request.args.get('count', '5'))

            pagination = CompanyModel.query.paginate(page, count, False)
            company = pagination.items

            if len(company) == 0:
                return self.response(404)

        else:
            company = CompanyModel.query.filter_by(id=company_id).first()
            if company is None:
                return self.response(404)

        return self.response(200, company, pagination)

    def post(self):
        """
        Create a new Company

        :return: Code 201
        """

        name = request.form.get('name')
        address = request.form.get('address')
        city = request.form.get('city')
        country = request.form.get('country')
        email = request.form.get('email')
        phone = request.form.get('phone')

        if name is None or address is None or city is None or country is None:
            self.response(400)

        company = CompanyModel(
            name=name,
            address=address,
            city=city,
            country=country,
            email=email,
            phone=phone
        )

        company.save()
        return self.response(201, {'id': company.id})

    def put(self, company_id):
        """
        Edit a Company

        :param company_id:
        :return: Code 200
        """

        if company_id is None:
            return self.response(400)

        company = CompanyModel.query.filter_by(id=company_id).first()
        if company is None:
            return self.response(404)

        name = request.form.get('name')
        if name:
            company.name = name

        address = request.form.get('address')
        if address:
            company.address = address

        city = request.form.get('city')
        if city:
            company.city = city

        country = request.form.get('country')
        if country:
            company.country = country

        email = request.form.get('email')
        if email:
            company.email = email

        phone = request.form.get('phone')
        if phone:
            company.phone = phone

        company.save()
        return self.response(200)

    def delete(self, company_id):
        """
        Delete a Company

        :param company_id:
        :return:
        """

        if company_id is None:
            return self.response(400)

        company = CompanyModel.query.filter_by(id=company_id).first()
        if company is None:
            return self.response(404)

        company.delete()
        return self.response(200)

companies_view = CompanyController.as_view('companies')
companies.add_url_rule('', view_func=companies_view, methods=['POST'])
companies.add_url_rule('', defaults={'company_id': None}, view_func=companies_view, methods=['GET'])
companies.add_url_rule('/<int:company_id>', view_func=companies_view, methods=['GET', 'PUT', 'DELETE'])<|endoftext|>


<commit_before>from pandas.core.base import PandasObject # noqa

try:
    from pandas.api.types import infer_dtype as infer_dtype # noqa
    _infer_dtype_kwds = {'skipna': False}
except ImportError:  # Pandas before 0.20.0
    from pandas.lib import infer_dtype as infer_dtype
    _infer_dtype_kwds = {}

from pandas import DataFrame, Series

try:
    # Import register decorators from pandas >= 0.23
    from pandas.api.extensions import (register_dataframe_accessor,
                                       register_series_accessor)
except ImportError:
    try:
        from pandas.core.accessor import AccessorProperty
    except ImportError:  # Pandas before 0.22.0
        from pandas.core.base import AccessorProperty

    # Define register decorators for pandas < 0.23
    class register_dataframe_accessor(object):
        """Register custom accessor on DataFrame."""

        def __init__(self, name):
            self.name = name

        def __call__(self, accessor):
            setattr(DataFrame, self.name, AccessorProperty(accessor, accessor))
            return accessor

    class register_series_accessor(object):
        """Register custom accessor on Series."""

        def __init__(self, name):
            self.name = name

        def __call__(self, accessor):
            setattr(Series, self.name, AccessorProperty(accessor, accessor))
            return accessor
<commit_msg>MAINT: fix flake8 failure<commit_after># flake8: noqa
from pandas.core.base import PandasObject

try:
    from pandas.api.types import infer_dtype as infer_dtype
    _infer_dtype_kwds = {'skipna': False}
except ImportError:  # Pandas before 0.20.0
    from pandas.lib import infer_dtype as infer_dtype
    _infer_dtype_kwds = {}

from pandas import DataFrame, Series

try:
    # Import register decorators from pandas >= 0.23
    from pandas.api.extensions import (register_dataframe_accessor,
                                       register_series_accessor)
except ImportError:
    try:
        from pandas.core.accessor import AccessorProperty
    except ImportError:  # Pandas before 0.22.0
        from pandas.core.base import AccessorProperty

    # Define register decorators for pandas < 0.23
    class register_dataframe_accessor(object):
        """Register custom accessor on DataFrame."""

        def __init__(self, name):
            self.name = name

        def __call__(self, accessor):
            setattr(DataFrame, self.name, AccessorProperty(accessor, accessor))
            return accessor

    class register_series_accessor(object):
        """Register custom accessor on Series."""

        def __init__(self, name):
            self.name = name

        def __call__(self, accessor):
            setattr(Series, self.name, AccessorProperty(accessor, accessor))
            return accessor
<|endoftext|>


<commit_before>import logging

from typing import Any, Dict, List, NoReturn, Optional, Text, Tuple, Type
from tqdm import tqdm

from rasa.nlu.tokenizers.convert_tokenizer import ConveRTTokenizer
from rasa.constants import DOCS_URL_COMPONENTS
from rasa.nlu.tokenizers.tokenizer import Token
from rasa.nlu.components import Component
from rasa.nlu.featurizers.featurizer import DenseFeaturizer, Features
from rasa.nlu.config import RasaNLUModelConfig
from rasa.nlu.training_data import Message, TrainingData
from rasa.nlu.constants import (
    TEXT,
    DENSE_FEATURIZABLE_ATTRIBUTES,
    FEATURIZER_CLASS_ALIAS,
    FEATURE_TYPE_SEQUENCE,
    FEATURE_TYPE_SENTENCE,
    TOKENS_NAMES,
)
import numpy as np
import tensorflow as tf

import rasa.utils.train_utils as train_utils
import rasa.utils.common as common_utils

logger = logging.getLogger(__name__)

class ConveRTFeaturizer(DenseFeaturizer):
    """Featurizer using ConveRT model.

    Loads the ConveRT(https://github.com/PolyAI-LDN/polyai-models#convert)
    model from TFHub and computes sentence and sequence level feature representations
    for dense featurizable attributes of each message object.
    """

    @classmethod
    def required_components(cls) -> List[Type[Component]]:
        return [ConveRTTokenizer]

    @classmethod
    def required_packages(cls) -> List[Text]:
        return ["tensorflow_text", "tensorflow_hub"]

    def __init__(self, component_config: Optional[Dict[Text, Any]] = None) -> None:

        super(ConveRTFeaturizer, self).__init__(component_config)

    def __get_signature(self, signature: Text, module: Any) -> NoReturn:
        """Retrieve a signature from a (hopefully loaded) TF model."""

        if not module:
            raise Exception(
                "ConveRTFeaturizer needs a proper loaded tensorflow module when used. "
                "Make sure to pass a module when training and using the component."
            )

        return module.signatures[signature]

    def _compute_features(
        self, batch_examples: List[Message], module: Any, attribute: Text = TEXT
    ) -> Tuple[np.ndarray, np.ndarray]:

        sentence_encodings = self._compute_sentence_encodings(
            batch_examples, module, attribute
        )

        (
            sequence_encodings,
            number_of_tokens_in_sentence,
        ) = self._compute_sequence_encodings(batch_examples, module, attribute)

        return self._get_features(
            sentence_encodings, sequence_encodings, number_of_tokens_in_sentence
        )

    def _compute_sentence_encodings(
        self, batch_examples: List[Message], module: Any, attribute: Text = TEXT
    ) -> np.ndarray:
        # Get text for attribute of each example
        batch_attribute_text = [ex.get(attribute) for ex in batch_examples]
        sentence_encodings = self._sentence_encoding_of_text(
            batch_attribute_text, module
        )

        # convert them to a sequence of 1
        return np.reshape(sentence_encodings, (len(batch_examples), 1, -1))

    def _compute_sequence_encodings(
        self, batch_examples: List[Message], module: Any, attribute: Text = TEXT
    ) -> Tuple[np.ndarray, List[int]]:
        list_of_tokens = [
            example.get(TOKENS_NAMES[attribute]) for example in batch_examples
        ]

        number_of_tokens_in_sentence = [
            len(sent_tokens) for sent_tokens in list_of_tokens
        ]

        # join the tokens to get a clean text to ensure the sequence length of
        # the returned embeddings from ConveRT matches the length of the tokens
        # (including sub-tokens)
        tokenized_texts = self._tokens_to_text(list_of_tokens)
        token_features = self._sequence_encoding_of_text(tokenized_texts, module)

        # ConveRT might split up tokens into sub-tokens
        # take the mean of the sub-token vectors and use that as the token vector
        token_features = train_utils.align_token_features(
            list_of_tokens, token_features
        )

        return token_features, number_of_tokens_in_sentence

    @staticmethod
    def _get_features(
        sentence_encodings: np.ndarray,
        sequence_encodings: np.ndarray,
        number_of_tokens_in_sentence: List[int],
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Get the sequence and sentence features."""

        sentence_embeddings = []
        sequence_embeddings = []

        for index in range(len(number_of_tokens_in_sentence)):
            sequence_length = number_of_tokens_in_sentence[index]
            sequence_encoding = sequence_encodings[index][:sequence_length]
            sentence_encoding = sentence_encodings[index]

            sequence_embeddings.append(sequence_encoding)
            sentence_embeddings.append(sentence_encoding)

        return np.array(sequence_embeddings), np.array(sentence_embeddings)

    @staticmethod
    def _tokens_to_text(list_of_tokens: List[List[Token]]) -> List[Text]:
        """Convert list of tokens to text.

        Add a whitespace between two tokens if the end value of the first tokens is
        not the same as the end value of the second token."""
        texts = []
        for tokens in list_of_tokens:
            text = ""
            offset = 0
            for token in tokens:
                if offset != token.start:
                    text += " "
                text += token.text

                offset = token.end
            texts.append(text)

        return texts

    def _sentence_encoding_of_text(self, batch: List[Text], module: Any) -> np.ndarray:
        signature = self.__get_signature("default", module)
        return signature(tf.convert_to_tensor(batch))["default"].numpy()

    def _sequence_encoding_of_text(self, batch: List[Text], module: Any) -> np.ndarray:
        signature = self.__get_signature("encode_sequence", module)

        return signature(tf.convert_to_tensor(batch))["sequence_encoding"].numpy()

    def train(
        self,
        training_data: TrainingData,
        config: Optional[RasaNLUModelConfig] = None,
        *,
        tf_hub_module: Any = None,
        **kwargs: Any,
    ) -> None:
        if config is not None and config.language != "en":
            common_utils.raise_warning(
                f"Since ``ConveRT`` model is trained only on an english "
                f"corpus of conversations, this featurizer should only be "
                f"used if your training data is in english language. "
                f"However, you are training in '{config.language}'. ",
                docs=DOCS_URL_COMPONENTS + "#convertfeaturizer",
            )

        batch_size = 64

        for attribute in DENSE_FEATURIZABLE_ATTRIBUTES:

            non_empty_examples = list(
                filter(lambda x: x.get(attribute), training_data.training_examples)
            )

            progress_bar = tqdm(
                range(0, len(non_empty_examples), batch_size),
                desc=attribute.capitalize() + " batches",
            )
            for batch_start_index in progress_bar:
                batch_end_index = min(
                    batch_start_index + batch_size, len(non_empty_examples)
                )

                # Collect batch examples
                batch_examples = non_empty_examples[batch_start_index:batch_end_index]

                (
                    batch_sequence_features,
                    batch_sentence_features,
                ) = self._compute_features(batch_examples, tf_hub_module, attribute)

                self._set_features(
                    batch_examples,
                    batch_sequence_features,
                    batch_sentence_features,
                    attribute,
                )

    def process(
        self, message: Message, *, tf_hub_module: Any = None, **kwargs: Any,
    ) -> None:

        for attribute in DENSE_FEATURIZABLE_ATTRIBUTES:
            sequence_features, sentence_features = self._compute_features(
                [message], tf_hub_module, attribute=attribute
            )

            self._set_features(
                [message], sequence_features, sentence_features, attribute
            )

    def _set_features(
        self,
        examples: List[Message],
        sequence_features: np.ndarray,
        sentence_features: np.ndarray,
        attribute: Text,
    ) -> None:
        for index, example in enumerate(examples):
            _sequence_features = Features(
                sequence_features[index],
                FEATURE_TYPE_SEQUENCE,
                attribute,
                self.component_config[FEATURIZER_CLASS_ALIAS],
            )
            example.add_features(_sequence_features)

            _sentence_features = Features(
                sentence_features[index],
                FEATURE_TYPE_SENTENCE,
                attribute,
                self.component_config[FEATURIZER_CLASS_ALIAS],
            )
            example.add_features(_sentence_features)
<commit_msg>ignore empty message<commit_after>import logging

from typing import Any, Dict, List, NoReturn, Optional, Text, Tuple, Type
from tqdm import tqdm

from rasa.nlu.tokenizers.convert_tokenizer import ConveRTTokenizer
from rasa.constants import DOCS_URL_COMPONENTS
from rasa.nlu.tokenizers.tokenizer import Token
from rasa.nlu.components import Component
from rasa.nlu.featurizers.featurizer import DenseFeaturizer, Features
from rasa.nlu.config import RasaNLUModelConfig
from rasa.nlu.training_data import Message, TrainingData
from rasa.nlu.constants import (
    TEXT,
    DENSE_FEATURIZABLE_ATTRIBUTES,
    FEATURIZER_CLASS_ALIAS,
    FEATURE_TYPE_SEQUENCE,
    FEATURE_TYPE_SENTENCE,
    TOKENS_NAMES,
)
import numpy as np
import tensorflow as tf

import rasa.utils.train_utils as train_utils
import rasa.utils.common as common_utils

logger = logging.getLogger(__name__)

class ConveRTFeaturizer(DenseFeaturizer):
    """Featurizer using ConveRT model.

    Loads the ConveRT(https://github.com/PolyAI-LDN/polyai-models#convert)
    model from TFHub and computes sentence and sequence level feature representations
    for dense featurizable attributes of each message object.
    """

    @classmethod
    def required_components(cls) -> List[Type[Component]]:
        return [ConveRTTokenizer]

    @classmethod
    def required_packages(cls) -> List[Text]:
        return ["tensorflow_text", "tensorflow_hub"]

    def __init__(self, component_config: Optional[Dict[Text, Any]] = None) -> None:

        super(ConveRTFeaturizer, self).__init__(component_config)

    def __get_signature(self, signature: Text, module: Any) -> NoReturn:
        """Retrieve a signature from a (hopefully loaded) TF model."""

        if not module:
            raise Exception(
                "ConveRTFeaturizer needs a proper loaded tensorflow module when used. "
                "Make sure to pass a module when training and using the component."
            )

        return module.signatures[signature]

    def _compute_features(
        self, batch_examples: List[Message], module: Any, attribute: Text = TEXT
    ) -> Tuple[np.ndarray, np.ndarray]:

        sentence_encodings = self._compute_sentence_encodings(
            batch_examples, module, attribute
        )

        (
            sequence_encodings,
            number_of_tokens_in_sentence,
        ) = self._compute_sequence_encodings(batch_examples, module, attribute)

        return self._get_features(
            sentence_encodings, sequence_encodings, number_of_tokens_in_sentence
        )

    def _compute_sentence_encodings(
        self, batch_examples: List[Message], module: Any, attribute: Text = TEXT
    ) -> np.ndarray:
        # Get text for attribute of each example
        batch_attribute_text = [ex.get(attribute) for ex in batch_examples]
        sentence_encodings = self._sentence_encoding_of_text(
            batch_attribute_text, module
        )

        # convert them to a sequence of 1
        return np.reshape(sentence_encodings, (len(batch_examples), 1, -1))

    def _compute_sequence_encodings(
        self, batch_examples: List[Message], module: Any, attribute: Text = TEXT
    ) -> Tuple[np.ndarray, List[int]]:
        list_of_tokens = [
            example.get(TOKENS_NAMES[attribute]) for example in batch_examples
        ]

        number_of_tokens_in_sentence = [
            len(sent_tokens) for sent_tokens in list_of_tokens
        ]

        # join the tokens to get a clean text to ensure the sequence length of
        # the returned embeddings from ConveRT matches the length of the tokens
        # (including sub-tokens)
        tokenized_texts = self._tokens_to_text(list_of_tokens)
        token_features = self._sequence_encoding_of_text(tokenized_texts, module)

        # ConveRT might split up tokens into sub-tokens
        # take the mean of the sub-token vectors and use that as the token vector
        token_features = train_utils.align_token_features(
            list_of_tokens, token_features
        )

        return token_features, number_of_tokens_in_sentence

    @staticmethod
    def _get_features(
        sentence_encodings: np.ndarray,
        sequence_encodings: np.ndarray,
        number_of_tokens_in_sentence: List[int],
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Get the sequence and sentence features."""

        sentence_embeddings = []
        sequence_embeddings = []

        for index in range(len(number_of_tokens_in_sentence)):
            sequence_length = number_of_tokens_in_sentence[index]
            sequence_encoding = sequence_encodings[index][:sequence_length]
            sentence_encoding = sentence_encodings[index]

            sequence_embeddings.append(sequence_encoding)
            sentence_embeddings.append(sentence_encoding)

        return np.array(sequence_embeddings), np.array(sentence_embeddings)

    @staticmethod
    def _tokens_to_text(list_of_tokens: List[List[Token]]) -> List[Text]:
        """Convert list of tokens to text.

        Add a whitespace between two tokens if the end value of the first tokens is
        not the same as the end value of the second token."""
        texts = []
        for tokens in list_of_tokens:
            text = ""
            offset = 0
            for token in tokens:
                if offset != token.start:
                    text += " "
                text += token.text

                offset = token.end
            texts.append(text)

        return texts

    def _sentence_encoding_of_text(self, batch: List[Text], module: Any) -> np.ndarray:
        signature = self.__get_signature("default", module)
        return signature(tf.convert_to_tensor(batch))["default"].numpy()

    def _sequence_encoding_of_text(self, batch: List[Text], module: Any) -> np.ndarray:
        signature = self.__get_signature("encode_sequence", module)

        return signature(tf.convert_to_tensor(batch))["sequence_encoding"].numpy()

    def train(
        self,
        training_data: TrainingData,
        config: Optional[RasaNLUModelConfig] = None,
        *,
        tf_hub_module: Any = None,
        **kwargs: Any,
    ) -> None:
        if config is not None and config.language != "en":
            common_utils.raise_warning(
                f"Since ``ConveRT`` model is trained only on an english "
                f"corpus of conversations, this featurizer should only be "
                f"used if your training data is in english language. "
                f"However, you are training in '{config.language}'. ",
                docs=DOCS_URL_COMPONENTS + "#convertfeaturizer",
            )

        batch_size = 64

        for attribute in DENSE_FEATURIZABLE_ATTRIBUTES:

            non_empty_examples = list(
                filter(lambda x: x.get(attribute), training_data.training_examples)
            )

            progress_bar = tqdm(
                range(0, len(non_empty_examples), batch_size),
                desc=attribute.capitalize() + " batches",
            )
            for batch_start_index in progress_bar:
                batch_end_index = min(
                    batch_start_index + batch_size, len(non_empty_examples)
                )

                # Collect batch examples
                batch_examples = non_empty_examples[batch_start_index:batch_end_index]

                (
                    batch_sequence_features,
                    batch_sentence_features,
                ) = self._compute_features(batch_examples, tf_hub_module, attribute)

                self._set_features(
                    batch_examples,
                    batch_sequence_features,
                    batch_sentence_features,
                    attribute,
                )

    def process(
        self, message: Message, *, tf_hub_module: Any = None, **kwargs: Any,
    ) -> None:

        for attribute in DENSE_FEATURIZABLE_ATTRIBUTES:
            if message.get(attribute):
                sequence_features, sentence_features = self._compute_features(
                    [message], tf_hub_module, attribute=attribute
                )

                self._set_features(
                    [message], sequence_features, sentence_features, attribute
                )

    def _set_features(
        self,
        examples: List[Message],
        sequence_features: np.ndarray,
        sentence_features: np.ndarray,
        attribute: Text,
    ) -> None:
        for index, example in enumerate(examples):
            _sequence_features = Features(
                sequence_features[index],
                FEATURE_TYPE_SEQUENCE,
                attribute,
                self.component_config[FEATURIZER_CLASS_ALIAS],
            )
            example.add_features(_sequence_features)

            _sentence_features = Features(
                sentence_features[index],
                FEATURE_TYPE_SENTENCE,
                attribute,
                self.component_config[FEATURIZER_CLASS_ALIAS],
            )
            example.add_features(_sentence_features)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Embed Flickr Images in Pelican Articles
=======================================

"""
import logging
import pickle
import re

import flickr as api_client

from pelican import signals

flickr_regex = re.compile(r'
(\[flickr:id\=([0-9]+)\])
')
default_template = """

    {{title}}

"""

logger = logging.getLogger(__name__)

def setup_flickr(generator):
    """Add Flickr api object to Pelican settings."""

    for key in ('TOKEN', 'KEY', 'SECRET'):
        try:
            value = generator.settings['FLICKR_API_' + key]
            setattr(api_client, 'API_' + key, value)
        except KeyError:
            logger.warning('[flickrtag]: FLICKR_API_%s is not defined in the configuration' % key)

    generator.flickr_api_client = api_client

    generator.settings.setdefault(
        'FLICKR_TAG_CACHE_LOCATION',
        '/tmp/com.chrisstreeter.flickrtag-images.cache')
    generator.settings.setdefault('FLICKR_TAG_INCLUDE_DIMENSIONS', False)
    generator.settings.setdefault('FLICKR_TAG_IMAGE_SIZE', 'Medium 640')

def url_for_alias(photo, alias):
    if alias == 'Medium 640':
        url = photo.getMedium640()
    else:
        url = photo.getMedium()
    url = url.replace('http:', '').replace('https:', '')
    return url

def size_for_alias(sizes, alias):
    if alias not in ('Medium 640', 'Medium'):
        alias = 'Medium'
    return [s for s in sizes if s['label'] == alias][0]

def generic_replace(generator, ct_type):
    if ct_type not in ('article', 'page'):
        ct_type = 'article'

    from jinja2 import Template

    api = generator.flickr_api_client
    if api is None:
        logger.error('[flickrtag]: Unable to get the Flickr API object')
        return

    tmp_file = generator.context.get('FLICKR_TAG_CACHE_LOCATION')

    include_dimensions = generator.context.get('FLICKR_TAG_INCLUDE_DIMENSIONS')
    size_alias = generator.context.get('FLICKR_TAG_IMAGE_SIZE')

    photo_ids = set([])
    logger.info('[flickrtag]: Parsing %ss for photo ids...' % ct_type)
    if ct_type == 'article':
        item_list = generator.articles
    else:
        item_list = generator.pages
    for item in item_list:
        for match in flickr_regex.findall(item._content):
            photo_ids.add(match[1])

    logger.info('[flickrtag]: Found %d photo ids in the %ss' % (len(photo_ids),ct_type) )

    try:
        with open(tmp_file, 'r') as f:
            photo_mapping = pickle.load(f)
    except (IOError, EOFError):
        photo_mapping = {}
    else:
        # Get the difference of photo_ids and what have cached
        cached_ids = set(photo_mapping.keys())
        photo_ids = list(set(photo_ids) - cached_ids)

    if photo_ids:
        logger.info('[flickrtag]: Fetching photo information from Flickr...')
        for id in photo_ids:
            logger.info('[flickrtag]: Fetching photo information for %s' % id)
            photo = api.Photo(id=id)
            # Trigger the API call...
            photo_mapping[id] = {
                'title': photo.title,
                'raw_url': url_for_alias(photo, size_alias),
                'url': photo.url,
            }

            if include_dimensions:
                sizes = photo.getSizes()
                size = size_for_alias(sizes, size_alias)
                photo_mapping[id]['width'] = size['width']
                photo_mapping[id]['height'] = size['height']

        with open(tmp_file, 'w') as f:
            pickle.dump(photo_mapping, f)
    else:
        logger.info('[flickrtag]: Found pickled photo mapping')

    # See if a custom template was provided
    template_name = generator.context.get('FLICKR_TAG_TEMPLATE_NAME')
    if template_name is not None:
        # There's a custom template
        try:
            template = generator.get_template(template_name)
        except Exception:
            logger.error('[flickrtag]: Unable to find the custom template %s' % template_name)
            template = Template(default_template)
    else:
        template = Template(default_template)

    logger.info('[flickrtag]: Inserting photo information into %ss...' % ct_type)
    if ct_type == 'article':
        item_list = generator.articles
    else:
        item_list = generator.pages
    for item in item_list:
        for match in flickr_regex.findall(item._content):
            fid = match[1]
            if fid not in photo_mapping:
                logger.error('[flickrtag]: Could not find info for a photo!')
                continue

            # Create a context to render with
            context = generator.context.copy()
            context.update(photo_mapping[fid])

            # Render the template
            replacement = template.render(context)

            item._content = item._content.replace(match[0], replacement)

def replace_article_tags(generator):
    generic_replace(generator, 'article')

def replace_page_tags(generator):
    generic_replace(generator, 'page')

def register():
    """Plugin registration."""
    signals.generator_init.connect(setup_flickr)

    signals.article_generator_finalized.connect(replace_article_tags)
    signals.page_generator_finalized.connect(replace_page_tags)
<commit_msg>new regex to catch multiple photo in same paragraph<commit_after># -*- coding: utf-8 -*-
"""
Embed Flickr Images in Pelican Articles
=======================================

"""
import logging
import pickle
import re

import flickr as api_client

from pelican import signals

flickr_regex = re.compile(r'(\[flickr:id\=([0-9]+)\])')
default_template = """

    {{title}}

"""

logger = logging.getLogger(__name__)

def setup_flickr(generator):
    """Add Flickr api object to Pelican settings."""

    for key in ('TOKEN', 'KEY', 'SECRET'):
        try:
            value = generator.settings['FLICKR_API_' + key]
            setattr(api_client, 'API_' + key, value)
        except KeyError:
            logger.warning('[flickrtag]: FLICKR_API_%s is not defined in the configuration' % key)

    generator.flickr_api_client = api_client

    generator.settings.setdefault(
        'FLICKR_TAG_CACHE_LOCATION',
        '/tmp/com.chrisstreeter.flickrtag-images.cache')
    generator.settings.setdefault('FLICKR_TAG_INCLUDE_DIMENSIONS', False)
    generator.settings.setdefault('FLICKR_TAG_IMAGE_SIZE', 'Medium 640')

def url_for_alias(photo, alias):
    if alias == 'Medium 640':
        url = photo.getMedium640()
    else:
        url = photo.getMedium()
    url = url.replace('http:', '').replace('https:', '')
    return url

def size_for_alias(sizes, alias):
    if alias not in ('Medium 640', 'Medium'):
        alias = 'Medium'
    return [s for s in sizes if s['label'] == alias][0]

def generic_replace(generator, ct_type):
    if ct_type not in ('article', 'page'):
        ct_type = 'article'

    from jinja2 import Template

    api = generator.flickr_api_client
    if api is None:
        logger.error('[flickrtag]: Unable to get the Flickr API object')
        return

    tmp_file = generator.context.get('FLICKR_TAG_CACHE_LOCATION')

    include_dimensions = generator.context.get('FLICKR_TAG_INCLUDE_DIMENSIONS')
    size_alias = generator.context.get('FLICKR_TAG_IMAGE_SIZE')

    photo_ids = set([])
    logger.info('[flickrtag]: Parsing %ss for photo ids...' % ct_type)
    if ct_type == 'article':
        item_list = generator.articles
    else:
        item_list = generator.pages
    for item in item_list:
        for match in flickr_regex.findall(item._content):
            photo_ids.add(match[1])

    logger.info('[flickrtag]: Found %d photo ids in the %ss' % (len(photo_ids),ct_type) )

    try:
        with open(tmp_file, 'r') as f:
            photo_mapping = pickle.load(f)
    except (IOError, EOFError):
        photo_mapping = {}
    else:
        # Get the difference of photo_ids and what have cached
        cached_ids = set(photo_mapping.keys())
        photo_ids = list(set(photo_ids) - cached_ids)

    if photo_ids:
        logger.info('[flickrtag]: Fetching photo information from Flickr...')
        for id in photo_ids:
            logger.info('[flickrtag]: Fetching photo information for %s' % id)
            photo = api.Photo(id=id)
            # Trigger the API call...
            photo_mapping[id] = {
                'title': photo.title,
                'raw_url': url_for_alias(photo, size_alias),
                'url': photo.url,
            }

            if include_dimensions:
                sizes = photo.getSizes()
                size = size_for_alias(sizes, size_alias)
                photo_mapping[id]['width'] = size['width']
                photo_mapping[id]['height'] = size['height']

        with open(tmp_file, 'w') as f:
            pickle.dump(photo_mapping, f)
    else:
        logger.info('[flickrtag]: Found pickled photo mapping')

    # See if a custom template was provided
    template_name = generator.context.get('FLICKR_TAG_TEMPLATE_NAME')
    if template_name is not None:
        # There's a custom template
        try:
            template = generator.get_template(template_name)
        except Exception:
            logger.error('[flickrtag]: Unable to find the custom template %s' % template_name)
            template = Template(default_template)
    else:
        template = Template(default_template)

    logger.info('[flickrtag]: Inserting photo information into %ss...' % ct_type)
    if ct_type == 'article':
        item_list = generator.articles
    else:
        item_list = generator.pages
    for item in item_list:
        for match in flickr_regex.findall(item._content):
            fid = match[1]
            if fid not in photo_mapping:
                logger.error('[flickrtag]: Could not find info for a photo!')
                continue

            # Create a context to render with
            context = generator.context.copy()
            context.update(photo_mapping[fid])

            # Render the template
            replacement = template.render(context)

            item._content = item._content.replace(match[0], replacement)

def replace_article_tags(generator):
    generic_replace(generator, 'article')

def replace_page_tags(generator):
    generic_replace(generator, 'page')

def register():
    """Plugin registration."""
    signals.generator_init.connect(setup_flickr)

    signals.article_generator_finalized.connect(replace_article_tags)
    signals.page_generator_finalized.connect(replace_page_tags)
<|endoftext|>