<commit_before>"""
Package for interacting on the network at a high level.

Copyright (c) 2014 Brian Muller
Copyright (c) 2015 OpenBazaar
"""

import pickle
import httplib
from binascii import hexlify
from twisted.internet.task import LoopingCall
from twisted.internet import defer, reactor, task

import nacl.signing
import nacl.hash
import nacl.encoding

from seed import peers
from log import Logger
from dht.protocol import KademliaProtocol
from dht.utils import deferredDict, digest
from dht.storage import ForgetfulStorage
from dht.node import Node
from dht.crawling import ValueSpiderCrawl
from dht.crawling import NodeSpiderCrawl

from protos import objects

def _anyRespondSuccess(responses):
    """
    Given the result of a DeferredList of calls to peers, ensure that at least
    one of them was contacted and responded with a Truthy result.
    """
    for deferSuccess, result in responses:
        peerReached, peerResponse = result
        if deferSuccess and peerReached and peerResponse:
            return True
    return False

class Server(object):
    """
    High level view of a node instance.  This is the object that should be created
    to start listening as an active node on the network.
    """

    def __init__(self, node, db, ksize=20, alpha=3, storage=None):
        """
        Create a server instance.  This will start listening on the given port.

        Args:
            node: The node instance for this peer. It must contain (at minimum) an ID,
                public key, ip address, and port.
            ksize (int): The k parameter from the paper
            alpha (int): The alpha parameter from the paper
            storage: An instance that implements :interface:`~dht.storage.IStorage`
        """
        self.ksize = ksize
        self.alpha = alpha
        self.log = Logger(system=self)
        self.storage = storage or ForgetfulStorage()
        self.node = node
        self.protocol = KademliaProtocol(self.node, self.storage, ksize, db)
        self.refreshLoop = LoopingCall(self.refreshTable).start(3600)

    def listen(self, port):
        """
        Start listening on the given port.

        This is the same as calling::

            reactor.listenUDP(port, server.protocol)
        """
        return reactor.listenUDP(port, self.protocol)

    def refreshTable(self):
        """
        Refresh buckets that haven't had any lookups in the last hour
        (per section 2.3 of the paper).
        """
        ds = []
        for rid in self.protocol.getRefreshIDs():
            node = Node(rid)
            nearest = self.protocol.router.findNeighbors(node, self.alpha)
            spider = NodeSpiderCrawl(self.protocol, node, nearest, self.ksize, self.alpha)
            ds.append(spider.find())

        def republishKeys(_):
            ds = []
            # Republish keys older than one hour
            for keyword in self.storage.iterkeys():
                for k, v in self.storage.iteritems(keyword):
                    if self.storage.get_ttl(keyword, k) < 601200:
                        ds.append(self.set(keyword, k, v))

        return defer.gatherResults(ds).addCallback(republishKeys)

    def querySeed(self, seed, pubkey):
        """
        Query an HTTP seed and return a `list` if (ip, port) `tuple` pairs.

        Args:
           seed: A `string` consisting of "ip:port" or "hostname:port"
           pubkey: The hex encoded public key to verify the signature on the response
        """
        try:
            nodes = []
            c = httplib.HTTPConnection(seed)
            c.request("GET", "/")
            response = c.getresponse()
            self.log.info("Https response from %s: %s, %s" % (seed, response.status, response.reason))
            data = response.read()
            reread_data = data.decode("zlib")
            proto = peers.PeerSeeds()
            proto.ParseFromString(reread_data)
            for peer in proto.peer_data:
                p = peers.PeerData()
                p.ParseFromString(peer)
                tup = (str(p.ip_address), p.port)
                nodes.append(tup)
            verify_key = nacl.signing.VerifyKey(pubkey, encoder=nacl.encoding.HexEncoder)
            verify_key.verify("".join(proto.peer_data), proto.signature)
            return nodes
        except Exception, e:
            self.log.error("Failed to query seed: %s" % str(e))

    def bootstrappableNeighbors(self):
        """
        Get a :class:`list` of (ip, port) :class:`tuple` pairs suitable for use as an argument
        to the bootstrap method.

        The server should have been bootstrapped
        already - this is just a utility for getting some neighbors and then
        storing them if this server is going down for a while.  When it comes
        back up, the list of nodes can be used to bootstrap.
        """
        neighbors = self.protocol.router.findNeighbors(self.node)
        return [tuple(n)[-2:] for n in neighbors]

    def bootstrap(self, addrs):
        """
        Bootstrap the server by connecting to other known nodes in the network.

        Args:
            addrs: A `list` of (ip, port) `tuple` pairs.  Note that only IP addresses
                   are acceptable - hostnames will cause an error.
        """

        # if the transport hasn't been initialized yet, wait a second
        if self.protocol.multiplexer.transport is None:
            return task.deferLater(reactor, 1, self.bootstrap, addrs)

        def initTable(results):
            nodes = []
            for addr, result in results.items():
                if result[0]:
                    n = objects.Node()
                    try:
                        n.ParseFromString(result[1][0])
                        pubkey = n.signedPublicKey[len(n.signedPublicKey) - 32:]
                        verify_key = nacl.signing.VerifyKey(pubkey)
                        verify_key.verify(n.signedPublicKey)
                        h = nacl.hash.sha512(n.signedPublicKey)
                        hash_pow = h[64:128]
                        if int(hash_pow[:6], 16) >= 50 or hexlify(n.guid) != h[:40]:
                            raise Exception('Invalid GUID')
                        nodes.append(Node(n.guid, addr[0], addr[1], n.signedPublicKey))
                    except Exception:
                        self.log.msg("Bootstrap node returned invalid GUID")
            spider = NodeSpiderCrawl(self.protocol, self.node, nodes, self.ksize, self.alpha)
            return spider.find()

        ds = {}
        for addr in addrs:
            ds[addr] = self.protocol.ping((addr[0], addr[1]))
        return deferredDict(ds).addCallback(initTable)

    def inetVisibleIP(self):
        """
        Get the internet visible IP's of this node as other nodes see it.

        Returns:
            A `list` of IP's.  If no one can be contacted, then the `list` will be empty.
        """

        def handle(results):
            ips = []
            for result in results:
                if result[0]:
                    ips.append((result[1][0], int(result[1][1])))
            self.log.debug("other nodes think our ip is %s" % str(ips))
            return ips

        ds = []
        for neighbor in self.bootstrappableNeighbors():
            ds.append(self.protocol.stun(neighbor))
        return defer.gatherResults(ds).addCallback(handle)

    def get(self, keyword):
        """
        Get a key if the network has it.

        Returns:
            :class:`None` if not found, the value otherwise.
        """
        dkey = digest(keyword)
        if self.storage.get(dkey) is not None:
            return defer.succeed(self.storage.get(dkey))
        node = Node(dkey)
        nearest = self.protocol.router.findNeighbors(node)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to get key %s" % keyword)
            return defer.succeed(None)
        spider = ValueSpiderCrawl(self.protocol, node, nearest, self.ksize, self.alpha)
        return spider.find()

    def set(self, keyword, key, value):
        """
        Set the given key/value tuple at the hash of the given keyword.
        All values stored in the DHT are stored as dictionaries of key/value
        pairs. If a value already exists for a given keyword, the new key/value
        pair will be appended to the dictionary.

        Args:
            keyword: The keyword to use. Should be hashed with hash160 before
                passing it in here.
            key: the 20 byte hash of the data.
            value: a serialized `protos.objects.Node` object which serves as a
                pointer to the node storing the data.

        Return: True if at least one peer responded. False if the store rpc
            completely failed.
        """
        if len(keyword) != 20:
            return defer.succeed(False)

        self.log.debug("setting '%s' = '%s':'%s' on network" % (keyword, hexlify(key), hexlify(value)))

        def store(nodes):
            self.log.info("setting '%s' on %s" % (keyword, [str(i) for i in nodes]))
            ds = [self.protocol.callStore(node, keyword, key, value) for node in nodes]

            keynode = Node(keyword)
            if self.node.distanceTo(keynode) < max([n.distanceTo(keynode) for n in nodes]):
                self.storage[keyword] = (key, value)
                self.log.debug("got a store request from %s, storing value" % str(self.node))

            return defer.DeferredList(ds).addCallback(_anyRespondSuccess)

        node = Node(keyword)
        nearest = self.protocol.router.findNeighbors(node)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to set key %s" % key)
            return defer.succeed(False)
        spider = NodeSpiderCrawl(self.protocol, node, nearest, self.ksize, self.alpha)
        return spider.find().addCallback(store)

    def delete(self, keyword, key, signature):
        """
        Delete the given key/value pair from the keyword dictionary on the network.
        To delete you must provide a signature covering the key that you wish to
        delete. It will be verified against the public key stored in the value. We
        use our ksize as alpha to make sure we reach as many nodes storing our value
        as possible.

        Args:
            keyword: the `string` keyword where the data being deleted is stored.
            key: the 20 byte hash of the data.
            signature: a signature covering the key.

        """
        self.log.info("deleting '%s':'%s' from the network" % (keyword, hexlify(key)))
        dkey = digest(keyword)

        def delete(nodes):
            self.log.debug("deleting '%s' on %s" % (key, [str(i) for i in nodes]))
            ds = [self.protocol.callDelete(node, dkey, key, signature) for node in nodes]

            if self.storage.getSpecific(dkey, key) is not None:
                self.storage.delete(dkey, key)

            return defer.DeferredList(ds).addCallback(_anyRespondSuccess)

        node = Node(dkey)
        nearest = self.protocol.router.findNeighbors(node)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to delete key %s" % key)
            return defer.succeed(False)
        spider = NodeSpiderCrawl(self.protocol, node, nearest, self.ksize, self.ksize)
        return spider.find().addCallback(delete)

    def resolve(self, guid):
        """
        Given a guid return a `Node` object containing its ip and port or none if it's
        not found.

        Args:
            guid: the 20 raw bytes representing the guid.
        """
        node_to_find = Node(guid)

        def check_for_node(nodes):
            for node in nodes:
                if node.id == node_to_find.id:
                    return node
            return None
        index = self.protocol.router.getBucketFor(node_to_find)
        nodes = self.protocol.router.buckets[index].getNodes()
        for node in nodes:
            if node.id == node_to_find.id:
                return defer.succeed(node)
        nearest = self.protocol.router.findNeighbors(node_to_find)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to find node %s" % node_to_find.id.encode("hex"))
            return defer.succeed(None)
        spider = NodeSpiderCrawl(self.protocol, node_to_find, nearest, self.ksize, self.alpha)
        return spider.find().addCallback(check_for_node)

    def saveState(self, fname):
        """
        Save the state of this node (the alpha/ksize/id/immediate neighbors)
        to a cache file with the given fname.
        """
        data = {'ksize': self.ksize,
                'alpha': self.alpha,
                'id': self.node.id,
                'vendor': self.node.vendor,
                'signed_pubkey': self.node.signed_pubkey,
                'neighbors': self.bootstrappableNeighbors(),
                'testnet': self.protocol.multiplexer.testnet}
        if len(data['neighbors']) == 0:
            self.log.warning("No known neighbors, so not writing to cache.")
            return
        with open(fname, 'w') as f:
            pickle.dump(data, f)

    @classmethod
    def loadState(cls, fname, ip_address, port, multiplexer, db, callback=None, storage=None):
        """
        Load the state of this node (the alpha/ksize/id/immediate neighbors)
        from a cache file with the given fname.
        """
        with open(fname, 'r') as f:
            data = pickle.load(f)
        if data['testnet'] != multiplexer.testnet:
            raise Exception('Cache uses wrong network parameters')
        n = Node(data['id'], ip_address, port, data['signed_pubkey'], data['vendor'])
        s = Server(n, db, data['ksize'], data['alpha'], storage=storage)
        s.protocol.connect_multiplexer(multiplexer)
        if len(data['neighbors']) > 0:
            if callback is not None:
                s.bootstrap(data['neighbors']).addCallback(callback)
            else:
                s.bootstrap(data['neighbors'])

        # TODO: if bootstrapping fails because all our cached nodes are dead, then query the seed again
        return s

    def saveStateRegularly(self, fname, frequency=600):
        """
        Save the state of node with a given regularity to the given
        filename.

        Args:
            fname: File name to save retularly to
            frequencey: Frequency in seconds that the state should be saved.
                        By default, 10 minutes.
        """
        loop = LoopingCall(self.saveState, fname)
        loop.start(frequency)
        return loop
<commit_msg>encode keys in log output<commit_after>"""
Package for interacting on the network at a high level.

Copyright (c) 2014 Brian Muller
Copyright (c) 2015 OpenBazaar
"""

import pickle
import httplib
from binascii import hexlify
from twisted.internet.task import LoopingCall
from twisted.internet import defer, reactor, task

import nacl.signing
import nacl.hash
import nacl.encoding

from seed import peers
from log import Logger
from dht.protocol import KademliaProtocol
from dht.utils import deferredDict, digest
from dht.storage import ForgetfulStorage
from dht.node import Node
from dht.crawling import ValueSpiderCrawl
from dht.crawling import NodeSpiderCrawl

from protos import objects

def _anyRespondSuccess(responses):
    """
    Given the result of a DeferredList of calls to peers, ensure that at least
    one of them was contacted and responded with a Truthy result.
    """
    for deferSuccess, result in responses:
        peerReached, peerResponse = result
        if deferSuccess and peerReached and peerResponse:
            return True
    return False

class Server(object):
    """
    High level view of a node instance.  This is the object that should be created
    to start listening as an active node on the network.
    """

    def __init__(self, node, db, ksize=20, alpha=3, storage=None):
        """
        Create a server instance.  This will start listening on the given port.

        Args:
            node: The node instance for this peer. It must contain (at minimum) an ID,
                public key, ip address, and port.
            ksize (int): The k parameter from the paper
            alpha (int): The alpha parameter from the paper
            storage: An instance that implements :interface:`~dht.storage.IStorage`
        """
        self.ksize = ksize
        self.alpha = alpha
        self.log = Logger(system=self)
        self.storage = storage or ForgetfulStorage()
        self.node = node
        self.protocol = KademliaProtocol(self.node, self.storage, ksize, db)
        self.refreshLoop = LoopingCall(self.refreshTable).start(3600)

    def listen(self, port):
        """
        Start listening on the given port.

        This is the same as calling::

            reactor.listenUDP(port, server.protocol)
        """
        return reactor.listenUDP(port, self.protocol)

    def refreshTable(self):
        """
        Refresh buckets that haven't had any lookups in the last hour
        (per section 2.3 of the paper).
        """
        ds = []
        for rid in self.protocol.getRefreshIDs():
            node = Node(rid)
            nearest = self.protocol.router.findNeighbors(node, self.alpha)
            spider = NodeSpiderCrawl(self.protocol, node, nearest, self.ksize, self.alpha)
            ds.append(spider.find())

        def republishKeys(_):
            ds = []
            # Republish keys older than one hour
            for keyword in self.storage.iterkeys():
                for k, v in self.storage.iteritems(keyword):
                    if self.storage.get_ttl(keyword, k) < 601200:
                        ds.append(self.set(keyword, k, v))

        return defer.gatherResults(ds).addCallback(republishKeys)

    def querySeed(self, seed, pubkey):
        """
        Query an HTTP seed and return a `list` if (ip, port) `tuple` pairs.

        Args:
           seed: A `string` consisting of "ip:port" or "hostname:port"
           pubkey: The hex encoded public key to verify the signature on the response
        """
        try:
            nodes = []
            c = httplib.HTTPConnection(seed)
            c.request("GET", "/")
            response = c.getresponse()
            self.log.info("Https response from %s: %s, %s" % (seed, response.status, response.reason))
            data = response.read()
            reread_data = data.decode("zlib")
            proto = peers.PeerSeeds()
            proto.ParseFromString(reread_data)
            for peer in proto.peer_data:
                p = peers.PeerData()
                p.ParseFromString(peer)
                tup = (str(p.ip_address), p.port)
                nodes.append(tup)
            verify_key = nacl.signing.VerifyKey(pubkey, encoder=nacl.encoding.HexEncoder)
            verify_key.verify("".join(proto.peer_data), proto.signature)
            return nodes
        except Exception, e:
            self.log.error("Failed to query seed: %s" % str(e))

    def bootstrappableNeighbors(self):
        """
        Get a :class:`list` of (ip, port) :class:`tuple` pairs suitable for use as an argument
        to the bootstrap method.

        The server should have been bootstrapped
        already - this is just a utility for getting some neighbors and then
        storing them if this server is going down for a while.  When it comes
        back up, the list of nodes can be used to bootstrap.
        """
        neighbors = self.protocol.router.findNeighbors(self.node)
        return [tuple(n)[-2:] for n in neighbors]

    def bootstrap(self, addrs):
        """
        Bootstrap the server by connecting to other known nodes in the network.

        Args:
            addrs: A `list` of (ip, port) `tuple` pairs.  Note that only IP addresses
                   are acceptable - hostnames will cause an error.
        """

        # if the transport hasn't been initialized yet, wait a second
        if self.protocol.multiplexer.transport is None:
            return task.deferLater(reactor, 1, self.bootstrap, addrs)

        def initTable(results):
            nodes = []
            for addr, result in results.items():
                if result[0]:
                    n = objects.Node()
                    try:
                        n.ParseFromString(result[1][0])
                        pubkey = n.signedPublicKey[len(n.signedPublicKey) - 32:]
                        verify_key = nacl.signing.VerifyKey(pubkey)
                        verify_key.verify(n.signedPublicKey)
                        h = nacl.hash.sha512(n.signedPublicKey)
                        hash_pow = h[64:128]
                        if int(hash_pow[:6], 16) >= 50 or hexlify(n.guid) != h[:40]:
                            raise Exception('Invalid GUID')
                        nodes.append(Node(n.guid, addr[0], addr[1], n.signedPublicKey))
                    except Exception:
                        self.log.msg("Bootstrap node returned invalid GUID")
            spider = NodeSpiderCrawl(self.protocol, self.node, nodes, self.ksize, self.alpha)
            return spider.find()

        ds = {}
        for addr in addrs:
            ds[addr] = self.protocol.ping((addr[0], addr[1]))
        return deferredDict(ds).addCallback(initTable)

    def inetVisibleIP(self):
        """
        Get the internet visible IP's of this node as other nodes see it.

        Returns:
            A `list` of IP's.  If no one can be contacted, then the `list` will be empty.
        """

        def handle(results):
            ips = []
            for result in results:
                if result[0]:
                    ips.append((result[1][0], int(result[1][1])))
            self.log.debug("other nodes think our ip is %s" % str(ips))
            return ips

        ds = []
        for neighbor in self.bootstrappableNeighbors():
            ds.append(self.protocol.stun(neighbor))
        return defer.gatherResults(ds).addCallback(handle)

    def get(self, keyword):
        """
        Get a key if the network has it.

        Returns:
            :class:`None` if not found, the value otherwise.
        """
        dkey = digest(keyword)
        if self.storage.get(dkey) is not None:
            return defer.succeed(self.storage.get(dkey))
        node = Node(dkey)
        nearest = self.protocol.router.findNeighbors(node)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to get key %s" % dkey.encode('hex'))
            return defer.succeed(None)
        spider = ValueSpiderCrawl(self.protocol, node, nearest, self.ksize, self.alpha)
        return spider.find()

    def set(self, keyword, key, value):
        """
        Set the given key/value tuple at the hash of the given keyword.
        All values stored in the DHT are stored as dictionaries of key/value
        pairs. If a value already exists for a given keyword, the new key/value
        pair will be appended to the dictionary.

        Args:
            keyword: The keyword to use. Should be hashed with hash160 before
                passing it in here.
            key: the 20 byte hash of the data.
            value: a serialized `protos.objects.Node` object which serves as a
                pointer to the node storing the data.

        Return: True if at least one peer responded. False if the store rpc
            completely failed.
        """
        if len(keyword) != 20:
            return defer.succeed(False)

        self.log.debug("setting '%s' on network" % keyword.encode("hex"))

        def store(nodes):
            self.log.info("setting '%s' on %s" % (keyword.encode("hex"), [str(i) for i in nodes]))
            ds = [self.protocol.callStore(node, keyword, key, value) for node in nodes]

            keynode = Node(keyword)
            if self.node.distanceTo(keynode) < max([n.distanceTo(keynode) for n in nodes]):
                self.storage[keyword] = (key, value)
                self.log.debug("got a store request from %s, storing value" % str(self.node))

            return defer.DeferredList(ds).addCallback(_anyRespondSuccess)

        node = Node(keyword)
        nearest = self.protocol.router.findNeighbors(node)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to set keyword %s" % keyword.encode("hex"))
            return defer.succeed(False)
        spider = NodeSpiderCrawl(self.protocol, node, nearest, self.ksize, self.alpha)
        return spider.find().addCallback(store)

    def delete(self, keyword, key, signature):
        """
        Delete the given key/value pair from the keyword dictionary on the network.
        To delete you must provide a signature covering the key that you wish to
        delete. It will be verified against the public key stored in the value. We
        use our ksize as alpha to make sure we reach as many nodes storing our value
        as possible.

        Args:
            keyword: the `string` keyword where the data being deleted is stored.
            key: the 20 byte hash of the data.
            signature: a signature covering the key.

        """
        self.log.info("deleting '%s':'%s' from the network" % (keyword.encode("hex"), key.encode("hex")))
        dkey = digest(keyword)

        def delete(nodes):
            self.log.debug("deleting '%s' on %s" % (key.encode("hex"), [str(i) for i in nodes]))
            ds = [self.protocol.callDelete(node, dkey, key, signature) for node in nodes]

            if self.storage.getSpecific(dkey, key) is not None:
                self.storage.delete(dkey, key)

            return defer.DeferredList(ds).addCallback(_anyRespondSuccess)

        node = Node(dkey)
        nearest = self.protocol.router.findNeighbors(node)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to delete key %s" % key.encode("hex"))
            return defer.succeed(False)
        spider = NodeSpiderCrawl(self.protocol, node, nearest, self.ksize, self.ksize)
        return spider.find().addCallback(delete)

    def resolve(self, guid):
        """
        Given a guid return a `Node` object containing its ip and port or none if it's
        not found.

        Args:
            guid: the 20 raw bytes representing the guid.
        """
        node_to_find = Node(guid)

        def check_for_node(nodes):
            for node in nodes:
                if node.id == node_to_find.id:
                    return node
            return None
        index = self.protocol.router.getBucketFor(node_to_find)
        nodes = self.protocol.router.buckets[index].getNodes()
        for node in nodes:
            if node.id == node_to_find.id:
                return defer.succeed(node)
        nearest = self.protocol.router.findNeighbors(node_to_find)
        if len(nearest) == 0:
            self.log.warning("There are no known neighbors to find node %s" % node_to_find.id.encode("hex"))
            return defer.succeed(None)
        spider = NodeSpiderCrawl(self.protocol, node_to_find, nearest, self.ksize, self.alpha)
        return spider.find().addCallback(check_for_node)

    def saveState(self, fname):
        """
        Save the state of this node (the alpha/ksize/id/immediate neighbors)
        to a cache file with the given fname.
        """
        data = {'ksize': self.ksize,
                'alpha': self.alpha,
                'id': self.node.id,
                'vendor': self.node.vendor,
                'signed_pubkey': self.node.signed_pubkey,
                'neighbors': self.bootstrappableNeighbors(),
                'testnet': self.protocol.multiplexer.testnet}
        if len(data['neighbors']) == 0:
            self.log.warning("No known neighbors, so not writing to cache.")
            return
        with open(fname, 'w') as f:
            pickle.dump(data, f)

    @classmethod
    def loadState(cls, fname, ip_address, port, multiplexer, db, callback=None, storage=None):
        """
        Load the state of this node (the alpha/ksize/id/immediate neighbors)
        from a cache file with the given fname.
        """
        with open(fname, 'r') as f:
            data = pickle.load(f)
        if data['testnet'] != multiplexer.testnet:
            raise Exception('Cache uses wrong network parameters')
        n = Node(data['id'], ip_address, port, data['signed_pubkey'], data['vendor'])
        s = Server(n, db, data['ksize'], data['alpha'], storage=storage)
        s.protocol.connect_multiplexer(multiplexer)
        if len(data['neighbors']) > 0:
            if callback is not None:
                s.bootstrap(data['neighbors']).addCallback(callback)
            else:
                s.bootstrap(data['neighbors'])

        # TODO: if bootstrapping fails because all our cached nodes are dead, then query the seed again
        return s

    def saveStateRegularly(self, fname, frequency=600):
        """
        Save the state of node with a given regularity to the given
        filename.

        Args:
            fname: File name to save retularly to
            frequencey: Frequency in seconds that the state should be saved.
                        By default, 10 minutes.
        """
        loop = LoopingCall(self.saveState, fname)
        loop.start(frequency)
        return loop
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# pylint: disable=duplicate-code

import copy

from datetime import datetime

import pytz

import singer
from singer import metadata
from singer import utils
from singer.schema import Schema

import pymysql.connections
import tap_mysql.sync_strategies.common as common

from tap_mysql.connection import make_connection_wrapper

from pymysqlreplication import BinLogStreamReader
from pymysqlreplication.event import RotateEvent
from pymysqlreplication.row_event import (
        DeleteRowsEvent,
        UpdateRowsEvent,
        WriteRowsEvent,
    )

LOGGER = singer.get_logger()

SDC_DELETED_AT = "_sdc_deleted_at"

UPDATE_BOOKMARK_PERIOD = 1000

def add_automatic_properties(catalog_entry):
    catalog_entry.schema.properties[SDC_DELETED_AT] = Schema(
        type=["null", "string"],
        format="date-time"
        )

def verify_binlog_config(connection, catalog_entry):
    with connection.cursor() as cur:
        cur.execute("""
        SELECT  @@binlog_format    AS binlog_format,
        @@binlog_row_image AS binlog_row_nimage;
        """)

        binlog_format, binlog_row_image = cur.fetchone()

        if binlog_format != 'ROW':
            raise Exception("""
            Unable to replicate with binlog for stream({}) because binlog_format is not set to 'ROW': {}.
            """.format(catalog_entry.stream, binlog_format))

        if binlog_row_image != 'FULL':
            raise Exception("""
            Unable to replicate with binlog for stream({}) because binlog_row_image is not set to 'FULL': {}.
            """.format(catalog_entry.stream, binlog_row_image))

def fetch_current_log_file_and_pos(connection):
    with connection.cursor() as cur:
        cur.execute("SHOW MASTER STATUS")

        result = cur.fetchone()
        current_log_file, current_log_pos = result[0:2]

        return current_log_file, current_log_pos

def fetch_server_id(connection):
    with connection.cursor() as cur:
        cur.execute("SELECT @@server_id")
        server_id = cur.fetchone()[0]

        return server_id

def row_to_singer_record(catalog_entry, version, vals, columns, time_extracted):
    filtered_vals = {k:v for k,v in vals.items() if k in columns}
    rec_cols, rec_vals = zip(*filtered_vals.items())

    return common.row_to_singer_record(catalog_entry, version, rec_vals, rec_cols, time_extracted)

def sync_table(connection, config, catalog_entry, state):
    verify_binlog_config(connection, catalog_entry)

    columns = common.generate_column_list(catalog_entry)

    if not columns:
        LOGGER.warning(
            'There are no columns selected for table %s, skipping it',
            catalog_entry.table)
        return

    stream_version = common.get_stream_version(catalog_entry.tap_stream_id, state)
    state = singer.write_bookmark(state,
                                  catalog_entry.tap_stream_id,
                                  'version',
                                  stream_version)

    yield singer.ActivateVersionMessage(
        stream=catalog_entry.stream,
        version=stream_version
    )

    server_id = fetch_server_id(connection)

    log_file = singer.get_bookmark(state,
                                   catalog_entry.tap_stream_id,
                                   'log_file')

    log_pos = singer.get_bookmark(state,
                                  catalog_entry.tap_stream_id,
                                  'log_pos')

    connection_wrapper = make_connection_wrapper(config)

    reader = BinLogStreamReader(
        connection_settings={},
        server_id=server_id,
        log_file=log_file,
        log_pos=log_pos,
        resume_stream=True,
        only_events=[RotateEvent, WriteRowsEvent, UpdateRowsEvent, DeleteRowsEvent],
        pymysql_wrapper=connection_wrapper
    )

    table_path = (catalog_entry.database, catalog_entry.stream)

    time_extracted = utils.now()

    LOGGER.info("Starting binlog replication with log_file=%s, log_pos=%s", log_file, log_pos)

    rows_saved = 0

    for binlog_event in reader:
        if isinstance(binlog_event, RotateEvent):
            #TODO
            print("ROTATE_EVENT")
        elif (binlog_event.schema, binlog_event.table) == table_path:
            if isinstance(binlog_event, WriteRowsEvent):
                for row in binlog_event.rows:
                    yield row_to_singer_record(catalog_entry,
                                               stream_version,
                                               row['values'],
                                               columns,
                                               time_extracted)
                    state = singer.write_bookmark(state,
                                                  catalog_entry.tap_stream_id,
                                                  'log_file',
                                                  log_file)
                    rows_saved = rows_saved + 1

            elif isinstance(binlog_event, UpdateRowsEvent):
                for row in binlog_event.rows:
                    yield row_to_singer_record(catalog_entry,
                                               stream_version,
                                               row['after_values'],
                                               columns,
                                               time_extracted)
                    rows_saved = rows_saved + 1
            elif isinstance(binlog_event, DeleteRowsEvent):
                for row in binlog_event.rows:
                    event_ts = datetime.utcfromtimestamp(binlog_event.timestamp).replace(tzinfo=pytz.UTC)

                    vals = row['values']
                    vals[SDC_DELETED_AT] = event_ts

                    yield row_to_singer_record(catalog_entry,
                                               stream_version,
                                               vals,
                                               columns,
                                               time_extracted)
                    rows_saved = rows_saved + 1

            state = singer.write_bookmark(state,
                                      catalog_entry.tap_stream_id,
                                      'log_file',
                                      reader.log_file)

            state = singer.write_bookmark(state,
                                      catalog_entry.tap_stream_id,
                                      'log_pos',
                                      reader.log_pos)

            if rows_saved % UPDATE_BOOKMARK_PERIOD == 0:
                yield singer.StateMessage(value=copy.deepcopy(state))

        yield singer.StateMessage(value=copy.deepcopy(state))
<commit_msg>do not process binlog event if log file and position match state<commit_after>#!/usr/bin/env python3
# pylint: disable=duplicate-code

import copy

from datetime import datetime

import pytz

import singer
from singer import metadata
from singer import utils
from singer.schema import Schema

import pymysql.connections
import tap_mysql.sync_strategies.common as common

from tap_mysql.connection import make_connection_wrapper

from pymysqlreplication import BinLogStreamReader
from pymysqlreplication.event import RotateEvent
from pymysqlreplication.row_event import (
        DeleteRowsEvent,
        UpdateRowsEvent,
        WriteRowsEvent,
    )

LOGGER = singer.get_logger()

SDC_DELETED_AT = "_sdc_deleted_at"

UPDATE_BOOKMARK_PERIOD = 1000

def add_automatic_properties(catalog_entry):
    catalog_entry.schema.properties[SDC_DELETED_AT] = Schema(
        type=["null", "string"],
        format="date-time"
        )

def verify_binlog_config(connection, catalog_entry):
    with connection.cursor() as cur:
        cur.execute("""
        SELECT  @@binlog_format    AS binlog_format,
        @@binlog_row_image AS binlog_row_nimage;
        """)

        binlog_format, binlog_row_image = cur.fetchone()

        if binlog_format != 'ROW':
            raise Exception("""
            Unable to replicate with binlog for stream({}) because binlog_format is not set to 'ROW': {}.
            """.format(catalog_entry.stream, binlog_format))

        if binlog_row_image != 'FULL':
            raise Exception("""
            Unable to replicate with binlog for stream({}) because binlog_row_image is not set to 'FULL': {}.
            """.format(catalog_entry.stream, binlog_row_image))

def fetch_current_log_file_and_pos(connection):
    with connection.cursor() as cur:
        cur.execute("SHOW MASTER STATUS")

        result = cur.fetchone()
        current_log_file, current_log_pos = result[0:2]

        return current_log_file, current_log_pos

def fetch_server_id(connection):
    with connection.cursor() as cur:
        cur.execute("SELECT @@server_id")
        server_id = cur.fetchone()[0]

        return server_id

def row_to_singer_record(catalog_entry, version, vals, columns, time_extracted):
    filtered_vals = {k:v for k,v in vals.items() if k in columns}
    rec_cols, rec_vals = zip(*filtered_vals.items())

    return common.row_to_singer_record(catalog_entry, version, rec_vals, rec_cols, time_extracted)

def sync_table(connection, config, catalog_entry, state):
    verify_binlog_config(connection, catalog_entry)

    columns = common.generate_column_list(catalog_entry)

    if not columns:
        LOGGER.warning(
            'There are no columns selected for table %s, skipping it',
            catalog_entry.table)
        return

    stream_version = common.get_stream_version(catalog_entry.tap_stream_id, state)
    state = singer.write_bookmark(state,
                                  catalog_entry.tap_stream_id,
                                  'version',
                                  stream_version)

    yield singer.ActivateVersionMessage(
        stream=catalog_entry.stream,
        version=stream_version
    )

    server_id = fetch_server_id(connection)

    log_file = singer.get_bookmark(state,
                                   catalog_entry.tap_stream_id,
                                   'log_file')

    log_pos = singer.get_bookmark(state,
                                  catalog_entry.tap_stream_id,
                                  'log_pos')

    connection_wrapper = make_connection_wrapper(config)

    reader = BinLogStreamReader(
        connection_settings={},
        server_id=server_id,
        log_file=log_file,
        log_pos=log_pos,
        resume_stream=True,
        only_events=[RotateEvent, WriteRowsEvent, UpdateRowsEvent, DeleteRowsEvent],
        pymysql_wrapper=connection_wrapper
    )

    table_path = (catalog_entry.database, catalog_entry.stream)

    time_extracted = utils.now()

    LOGGER.info("Starting binlog replication with log_file=%s, log_pos=%s", log_file, log_pos)

    rows_saved = 0

    for binlog_event in reader:
        if reader.log_file == log_file and reader.log_pos == log_pos:
            LOGGER.info("Skipping event for log_file=%s and log_pos=%s as it was processed last sync",
                        reader.log_file,
                        reader.log_pos)
            continue

        if isinstance(binlog_event, RotateEvent):
            #TODO
            print("ROTATE_EVENT")
        elif (binlog_event.schema, binlog_event.table) == table_path:
            if isinstance(binlog_event, WriteRowsEvent):
                for row in binlog_event.rows:
                    yield row_to_singer_record(catalog_entry,
                                               stream_version,
                                               row['values'],
                                               columns,
                                               time_extracted)
                    state = singer.write_bookmark(state,
                                                  catalog_entry.tap_stream_id,
                                                  'log_file',
                                                  log_file)
                    rows_saved = rows_saved + 1

            elif isinstance(binlog_event, UpdateRowsEvent):
                for row in binlog_event.rows:
                    yield row_to_singer_record(catalog_entry,
                                               stream_version,
                                               row['after_values'],
                                               columns,
                                               time_extracted)
                    rows_saved = rows_saved + 1
            elif isinstance(binlog_event, DeleteRowsEvent):
                for row in binlog_event.rows:
                    event_ts = datetime.utcfromtimestamp(binlog_event.timestamp).replace(tzinfo=pytz.UTC)

                    vals = row['values']
                    vals[SDC_DELETED_AT] = event_ts

                    yield row_to_singer_record(catalog_entry,
                                               stream_version,
                                               vals,
                                               columns,
                                               time_extracted)
                    rows_saved = rows_saved + 1

            state = singer.write_bookmark(state,
                                      catalog_entry.tap_stream_id,
                                      'log_file',
                                      reader.log_file)

            state = singer.write_bookmark(state,
                                      catalog_entry.tap_stream_id,
                                      'log_pos',
                                      reader.log_pos)

            if rows_saved % UPDATE_BOOKMARK_PERIOD == 0:
                yield singer.StateMessage(value=copy.deepcopy(state))

        yield singer.StateMessage(value=copy.deepcopy(state))
<|endoftext|>


<commit_before><commit_msg>update max year for RAIS<commit_after><|endoftext|>


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

# Copyright (c) 2014 Anton Tsyganenko
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os
import sys
import optparse
import hashlib
import base64

notenoughkeys = 5 # for warning (see KEY INPUT)

################ FUNCTIONS

def bxor(b1, b2): # use xor for bytes
    result = bytearray()
    for b1, b2 in zip(b1, b2):
        result.append(b1 ^ b2)
    return bytes(result)

def out(output):
    if outfile: # output to a file
        with open(outfile, "bw") as file:
            file.write(output)
    else:
        print(output.decode("utf-8"))

################ OPTIONS

parser = optparse.OptionParser()

parser.add_option("-o", "--output-file",
                  dest = "outfile",
                  help = "write output to FILE",
                  metavar = "FILE")

parser.add_option("-i", "--input-file",
                  dest = "infile",
                  help = "write output to FILE",
                  metavar = "FILE")

parser.add_option("-k", "--keys-folder",
                  dest = "keysource",
                  help = "get a key from FOLDER",
                  metavar = "FOLDER")

parser.add_option("-K", "--key-action",
                  dest = "keyaction",
                  choices = ["leave", "delete", "rename"],
                  default = "rename",
                  help = "action to do with used key, can be " +
                       "`leave`, `delete` or `rename`",
                  metavar = "ACTION")

parser.add_option("-I", "--input-mode",
                  dest = "imode",
                  default = "auto",
                  help = "input mode, can be `bin` or `b64`",
                  choices = ["bin", "b64", "auto"],
                  metavar = "MODE")

parser.add_option("-O", "--output-mode",
                  dest = "omode",
                  default = "auto",
                  choices = ["bin", "b64", "auto"],
                  help = "output mode, can be `bin` or `b64`",
                  metavar = "MODE")

parser.add_option("-g", "--gen-keys",
                  dest = "genkey",
                  default = False,
                  action = "store_true",
                  help = "generate keys")

parser.add_option("-c", "--hash-action",
                  dest = "hashaction",
                  choices = ["no", "auto", "add", "check"],
                  default = "auto",
                  metavar = "ACTION",
                  help = "action to do with hash sum")

(options, args) = parser.parse_args()

infile = options.infile
outfile = options.outfile
keysource = options.keysource
keyaction = options.keyaction
imode = options.imode
omode = options.omode
genkey = options.genkey
hashaction = options.hashaction

if hashaction in ["auto", "add"]:
    hashlen = 20
else:
    hashlen = 0

################ KEY GENERATION

if genkey:
    number = int(input("number of keys > "))
    length = int(input("key length > "))

    if omode == "b64" or not outfile: # text with keys in base64
        result = b""
        for i in range(number):
            result += base64.b64encode(os.urandom(length)) + b"\n"
        out(result)

    else: # folder with binary key files
        os.makedirs(outfile)
        for i in range(number):
            with open(outfile + os.sep + str(i), "xb") as f:
                f.write(os.urandom(length))

    exit()

################ TEXT INPUT

if infile: # from a file
    with open(infile, "rb") as file:
        text = file.read()

else: # direct input
    text = bytes(input("enter the text > "), "utf-8")

################ SETTINGS GUESSING

if infile:
    if imode == "auto":
        imode = "bin"

if imode in ["auto", "b64"]:
    try: # try to decode base64
        text = base64.b64decode(text.decode("utf-8"), validate=True)
    except:
        pass

# there are some other settings guessing code in FINAL.

################ KEY INPUT

if keysource: # use folder with keys or a keyfile
    if os.path.isdir(keysource): # folder with keys
        fileslist = os.listdir(keysource)

        for i in fileslist[:]: # don't use used keys
            if "_used" in i:
                fileslist.remove(i)

        if len(fileslist) == 0:
            print("================\n" +
                  "NO KEYS IN {kf}!".format(kf=keysource))
            exit()

        if len(fileslist) <= notenoughkeys:
            print("================")
            print(("WARNING! only {k} keys left, and one of them " +
                  "will be used now.").format(k=len(fileslist)))

        keyfile = keysource + os.sep + max(fileslist, key=int)
        with open(keyfile, "br") as f:
            key = f.read(len(text)+hashlen)

        keyfromfolder = True # see FINAL

        del fileslist
    else: # a keyfile
        with open(keysource, "br") as f:
            key = f.read(len(text)+hashlen)
        keyfromfolder = False

    del keysource, hashlen

else: # manually input the key
    key = base64.b64decode(input("enter the key > "))
    keyfromfolder = False

################ ADD A HASH SUM

if hashaction in ["add", "auto"]:
    text += hashlib.sha1(text).digest()

################ DON'T LET USER TO USE TOO SHORT KEY

if len(key) < (len(text)):
    print("The key must have the text's length or be longer")
    if hashaction == "auto" and len(key) >= (len(text)-20):
        print("If you're decrypting data, try to use `-c check` option")
    exit()

################ ENCRYPTION/DECRYPTION

result = bxor(text, key)
del text, key

################ CHECK A HASH SUM

if hashaction != "no":
    if hashaction == "check":
        hashplace = memoryview(result)[-20:]
        resultbody = memoryview(result)[0:-20]
    elif hashaction == "auto":
        # in auto mode the program adds new hash sum(see ADD A HASH SUM)
        hashplace = memoryview(result)[-40:-20]
        resultbody = memoryview(result)[0:-40]

    if hashaction != "add":
        resulthash = hashlib.sha1(resultbody).digest()

    print("================")

    if hashaction != "add" and resulthash == hashplace:
        print("The hash sum is ok")
        result = bytes(resultbody)

    elif hashaction == "check":
        print("The hash sum is wrong!")
        result = bytes(resultbody)

    elif hashaction == "add":
        print("A hash sum has been added")
    else:
        print("A hash sum was added/wrong")

    if hashaction != "add":
        del resultbody, hashplace

################ FINAL

if omode == "auto": # settings guessing
    if not outfile:
        try:
            result.decode("utf-8")
            omode = "bin"
        except:
            omode = "b64"
    else:
        omode = "bin"

if omode == "b64": # convert encrypted result to base64
    result = base64.b64encode(result)

if not outfile and not (infile and keysource and not hashaction):
    print("================\n") # separator

out(result)

if keyfromfolder: # delete or rename used key
    if keyaction == "rename":
        os.rename(keyfile, keyfile + "_used")
    elif keyaction == "delete":
        os.remove(keyfile)
<commit_msg>fixed condition to print separator<commit_after>#!/usr/bin/env python3

# Copyright (c) 2014 Anton Tsyganenko
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os
import sys
import optparse
import hashlib
import base64

notenoughkeys = 5 # for warning (see KEY INPUT)

################ FUNCTIONS

def bxor(b1, b2): # use xor for bytes
    result = bytearray()
    for b1, b2 in zip(b1, b2):
        result.append(b1 ^ b2)
    return bytes(result)

def out(output):
    if outfile: # output to a file
        with open(outfile, "bw") as file:
            file.write(output)
    else:
        print(output.decode("utf-8"))

################ OPTIONS

parser = optparse.OptionParser()

parser.add_option("-o", "--output-file",
                  dest = "outfile",
                  help = "write output to FILE",
                  metavar = "FILE")

parser.add_option("-i", "--input-file",
                  dest = "infile",
                  help = "write output to FILE",
                  metavar = "FILE")

parser.add_option("-k", "--keys-folder",
                  dest = "keysource",
                  help = "get a key from FOLDER",
                  metavar = "FOLDER")

parser.add_option("-K", "--key-action",
                  dest = "keyaction",
                  choices = ["leave", "delete", "rename"],
                  default = "rename",
                  help = "action to do with used key, can be " +
                       "`leave`, `delete` or `rename`",
                  metavar = "ACTION")

parser.add_option("-I", "--input-mode",
                  dest = "imode",
                  default = "auto",
                  help = "input mode, can be `bin` or `b64`",
                  choices = ["bin", "b64", "auto"],
                  metavar = "MODE")

parser.add_option("-O", "--output-mode",
                  dest = "omode",
                  default = "auto",
                  choices = ["bin", "b64", "auto"],
                  help = "output mode, can be `bin` or `b64`",
                  metavar = "MODE")

parser.add_option("-g", "--gen-keys",
                  dest = "genkey",
                  default = False,
                  action = "store_true",
                  help = "generate keys")

parser.add_option("-c", "--hash-action",
                  dest = "hashaction",
                  choices = ["no", "auto", "add", "check"],
                  default = "auto",
                  metavar = "ACTION",
                  help = "action to do with hash sum")

(options, args) = parser.parse_args()

infile = options.infile
outfile = options.outfile
keysource = options.keysource
keyaction = options.keyaction
imode = options.imode
omode = options.omode
genkey = options.genkey
hashaction = options.hashaction

if hashaction in ["auto", "add"]:
    hashlen = 20
else:
    hashlen = 0

################ KEY GENERATION

if genkey:
    number = int(input("number of keys > "))
    length = int(input("key length > "))

    if omode == "b64" or not outfile: # text with keys in base64
        result = b""
        for i in range(number):
            result += base64.b64encode(os.urandom(length)) + b"\n"
        out(result)

    else: # folder with binary key files
        os.makedirs(outfile)
        for i in range(number):
            with open(outfile + os.sep + str(i), "xb") as f:
                f.write(os.urandom(length))

    exit()

################ TEXT INPUT

if infile: # from a file
    with open(infile, "rb") as file:
        text = file.read()

else: # direct input
    text = bytes(input("enter the text > "), "utf-8")

################ SETTINGS GUESSING

if infile:
    if imode == "auto":
        imode = "bin"

if imode in ["auto", "b64"]:
    try: # try to decode base64
        text = base64.b64decode(text.decode("utf-8"), validate=True)
    except:
        pass

# there are some other settings guessing code in FINAL.

################ KEY INPUT

if keysource: # use folder with keys or a keyfile
    if os.path.isdir(keysource): # folder with keys
        fileslist = os.listdir(keysource)

        for i in fileslist[:]: # don't use used keys
            if "_used" in i:
                fileslist.remove(i)

        if len(fileslist) == 0:
            print("================\n" +
                  "NO KEYS IN {kf}!".format(kf=keysource))
            exit()

        if len(fileslist) <= notenoughkeys:
            print("================")
            print(("WARNING! only {k} keys left, and one of them " +
                  "will be used now.").format(k=len(fileslist)))

        keyfile = keysource + os.sep + max(fileslist, key=int)
        with open(keyfile, "br") as f:
            key = f.read(len(text)+hashlen)

        keyfromfolder = True # see FINAL

        del fileslist
    else: # a keyfile
        with open(keysource, "br") as f:
            key = f.read(len(text)+hashlen)
        keyfromfolder = False

    del keysource, hashlen

else: # manually input the key
    key = base64.b64decode(input("enter the key > "))
    keyfromfolder = False

################ ADD A HASH SUM

if hashaction in ["add", "auto"]:
    text += hashlib.sha1(text).digest()

################ DON'T LET USER TO USE TOO SHORT KEY

if len(key) < (len(text)):
    print("The key must have the text's length or be longer")
    if hashaction == "auto" and len(key) >= (len(text)-20):
        print("If you're decrypting data, try to use `-c check` option")
    exit()

################ ENCRYPTION/DECRYPTION

result = bxor(text, key)
del text, key

################ CHECK A HASH SUM

if hashaction != "no":
    if hashaction == "check":
        hashplace = memoryview(result)[-20:]
        resultbody = memoryview(result)[0:-20]
    elif hashaction == "auto":
        # in auto mode the program adds new hash sum(see ADD A HASH SUM)
        hashplace = memoryview(result)[-40:-20]
        resultbody = memoryview(result)[0:-40]

    if hashaction != "add":
        resulthash = hashlib.sha1(resultbody).digest()

    print("================")

    if hashaction != "add" and resulthash == hashplace:
        print("The hash sum is ok")
        result = bytes(resultbody)

    elif hashaction == "check":
        print("The hash sum is wrong!")
        result = bytes(resultbody)

    elif hashaction == "add":
        print("A hash sum has been added")
    else:
        print("A hash sum was added/wrong")

    if hashaction != "add":
        del resultbody, hashplace

################ FINAL

if omode == "auto": # settings guessing
    if not outfile:
        try:
            result.decode("utf-8")
            omode = "bin"
        except:
            omode = "b64"
    else:
        omode = "bin"

if omode == "b64": # convert encrypted result to base64
    result = base64.b64encode(result)

if not (outfile or infile and keyfromfolder and hashaction == "no"):
    print("================\n") # separator

out(result)

if keyfromfolder: # delete or rename used key
    if keyaction == "rename":
        os.rename(keyfile, keyfile + "_used")
    elif keyaction == "delete":
        os.remove(keyfile)
<|endoftext|>


<commit_before>"""Models for stories and story sections"""

from django.conf import settings
from django.contrib.auth.models import User
from django.contrib.sites.models import Site
from django.core import urlresolvers
from django.db import models
from django.db.models.signals import post_save, pre_save
from django.utils import simplejson
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext as _
from django_dag.models import edge_factory, node_factory
# TODO: Decide on tagging suggestion admin app.
# Right now, I'm using a hacked version of
# https://bitbucket.org/fabian/django-taggit-autosuggest
# which I modified to allow for specifying a Tag model
# other than taggit.models.Tag
#from taggit_autosuggest.managers import TaggableManager
from uuidfield.fields import UUIDField
from storybase.fields import ShortTextField
from storybase.models import (LicensedModel, PublishedModel,
    TimestampedModel, TranslatedModel, TranslationModel,
    set_date_on_published)
from storybase.utils import slugify
from storybase_asset.models import Asset
from storybase_user.models import Organization, Project
from storybase_story.managers import StoryManager
from storybase_story import structure
#from storybase_tag.models import TaggedItem

class StoryTranslation(TranslationModel):
    """Encapsulates translated fields of a Story"""
    story = models.ForeignKey('Story')
    title = ShortTextField()
    summary = models.TextField(blank=True)

    class Meta:
        """Model metadata options"""
        unique_together = (('story', 'language'))

    def __unicode__(self):
        return self.title

class Story(TranslatedModel, LicensedModel, PublishedModel,
            TimestampedModel):
    """Metadata for a story

    The Story model stores a story's metadata and aggregates a story's
    media assets

    """
    story_id = UUIDField(auto=True)
    slug = models.SlugField(blank=True)
    byline = models.TextField()
    structure = models.CharField(max_length=20,
        choices=structure.manager.get_structure_options())
    # blank=True, null=True to bypass validation so the user doesn't
    # have to always remember to set this in the Django admin.
    # Though this is set to blank=True, null=True, we should always set
    # this value.  In fact, the StoryModelAdmin class sets this to
    # request.user
    author = models.ForeignKey(User, related_name="stories", blank=True,
                               null=True)
    assets = models.ManyToManyField(Asset, related_name='stories',
                                    blank=True)
    featured_assets = models.ManyToManyField(Asset,
       related_name='featured_in_stories', blank=True,
       help_text=_("Assets to be displayed in teaser version of Story"))
    organizations = models.ManyToManyField(Organization,
                                           related_name='stories',
                                           blank=True)
    projects = models.ManyToManyField(Project, related_name='stories',
                                      blank=True)
    on_homepage = models.BooleanField(_("Featured on homepage"),
		                      default=False)
    #tags = TaggableManager(through=TaggedItem, blank=True)

    objects = StoryManager()

    translated_fields = ['title', 'summary']
    translation_set = 'storytranslation_set'

    _structure_obj = None

    class Meta:
        """Model metadata options"""
        verbose_name_plural = "stories"

    def __unicode__(self):
        return self.title

    @models.permalink
    def get_absolute_url(self):
        """Calculate the canonical URL for a Story"""
        return ('story_detail', [self.slug])

    def get_full_url(self):
        """
        Calculate the canonical URL for a Story, including its
        domain name.
        """
        return "http://%s%s" % (Site.objects.get_current().domain,
                                self.get_absolute_url())

    @property
    def contributor_name(self):
        """
        Return the contributor's first name and last initial or username
        """
        contributor = self.author
        contributor_name = contributor.username
        if contributor.first_name:
            contributor_name = contributor.first_name

            if contributor.last_name:
                contributor_name = "%s %s." % (contributor.first_name,
                                               contributor.last_name[0])
            else:
                contributor_name = contributor.first_name

        return contributor_name

    def get_structure_obj(self):
        """Return a structure object for the story"""
        if self._structure_obj is None:
            # A structure object hasn't been instantiated yet.
            # Create one.
            structure_class = structure.manager.get_structure_class(
            self.structure)
            self._structure_obj = structure_class.__call__(story=self)

        return self._structure_obj

    def sections_flat(self):
        """
        Return a list of sections ordered with each branch coming before
        the next
        """
        flattened = []
        for section in self.sections.filter(root=True).order_by('weight'):
            flattened.append(section)
            flattened = flattened + section.children_flat()

        return flattened

    def sections_json(self):
        """Return a JSON representation of the story sections"""
        sections = []
        for section in self.sections_flat():
            sections.append(section.to_simple())

        return mark_safe(simplejson.dumps(sections))

    def to_simple(self):
        """
        Return a simplified version of the story object that can be
        passed to json.dumps()
        """
        return {
            'story_id': self.story_id,
            'title': self.title,
            'summary': self.summary
        }

    def to_json(self):
        """Return JSON representation of this object"""
        return mark_safe(simplejson.dumps(self.to_simple()))

    def render_featured_asset(self, format='html'):
        """Render a representation of the story's featured asset"""
        try:
            # Return the first featured asset.  We have the ability of
            # selecting multiple featured assets.  Perhaps in the future
            # allow for specifying a particular feature asset or randomly
            # displaying one.
            featured_asset = self.featured_assets.select_subclasses()[0]
            thumbnail_options = {}
            if format == 'html':
                thumbnail_options.update({'html_class': 'featured-asset'})
            return featured_asset.render_thumbnail(format=format,
                                                   **thumbnail_options)
        except IndexError:
            # No featured assets
            return ''

    def featured_asset_thumbnail_url(self):
        """Return the URL of the featured asset's thumbnail

        Returns None if the asset cannot be converted to a thumbnail image.

        """
        try:
            featured_asset = self.featured_assets.select_subclasses()[0]
            return featured_asset.get_thumbnail_url()
        except IndexError:
            # No featured assets
            return None

    def render_story_structure(self, format='html'):
        """Render a representation of the Story structure"""
        output = []
        try:
            root = self.get_root_section()
        except Section.DoesNotExist:
            return ''

        output.append('
')
        output.append(root.render(format))
        output.append('
')
        return mark_safe(u'\n'.join(output))

    def render_toc(self, format='html'):
        """Render a representation of the Story's table of contents"""
        structure = self.get_structure_obj()
        return structure.render_toc(format)

def set_story_slug(sender, instance, **kwargs):
    """
    When a StoryTranslation is saved, set its Story's slug if it doesn't have
    one

    Should be connected to StoryTranslation's post_save signal.
    """
    if not instance.story.slug:
        instance.story.slug = slugify(instance.title)
	instance.story.save()

# Hook up some signal handlers
pre_save.connect(set_date_on_published, sender=Story)
post_save.connect(set_story_slug, sender=StoryTranslation)

class Section(node_factory('SectionRelation'), TranslatedModel):
    """ Section of a story """
    section_id = UUIDField(auto=True)
    story = models.ForeignKey('Story', related_name='sections')
    # True if this section the root section of the story, either
    # the first section in a linear story, or the central node
    # in a drill-down/"spider" structure.  Otherwise, False
    root = models.BooleanField(default=False)
    weight = models.IntegerField(default=0)
    """The ordering of top-level sections relative to each other"""
    assets = models.ManyToManyField(Asset, related_name='sections',
                                    blank=True, through='SectionAsset')

    translated_fields = ['title']
    translation_set = 'sectiontranslation_set'

    def __unicode__(self):
        try:
            return self.title
        except IndexError:
            # HACK: Need this to support deleting Sections through
            # an inline on the Story Change page
            #
            # When deleting an object in the Django admin, a Section
            # object gets instantiated with no translation set.
            # So, the call to __getattr__() raises an IndexError
            return super(Section, self).__unicode__()

    @property
    def sections_in_same_story(self):
        """Return other sections in the same story"""
        return self.__class__.objects.filter(story=self.story)

    def _child_relations(self):
        """
        Get a query set of through model instances for children
        of this section
        """
        return self.children.through.objects.filter(parent=self).order_by(
            'weight', 'child__sectiontranslation__title')

    def _parent_relations(self):
        """
        Get a list of through model instances for parents of this
        section
        """
        # This is a little confusing because we want parents, not
        # children.  self.children.through just gets us the "through" class
        # for the relation.
        return self.children.through.objects.filter(child=self).order_by(
            'weight', 'child__sectiontranslation__title')

    # TODO: Think about the best strategy for building a sense of
    # previous and next sections. I'm not sure if this is the best
    # place for the functionality and there's likely a way to avoid
    # having to traverse the tree so many times.

    def _next_child_section(self):
        """Return the first section that is a child of this section"""
        # We get it through the "through" class so we can order
        # the children by weight
        relation = self._child_relations()[0]
        return relation.child

    def _next_root_section(self):
        """Return the next top-level section"""
        if self.is_root() or self.is_island():
            root = self
        else:
            root = self.get_roots()[0]

        found = False
        for section in self.sections_in_same_story.order_by('weight'):
            if found:
                if section.is_root() or section.is_island():
                    return section
            else:
                if section == root:
                    found = True

        return None

    def get_next_section(self):
        """Get the next section"""
        # TODO: Unit test for this
        if self.children.count():
            # The section has children, return the first child
            section = self._next_child_section()
        else:
            # The section is a leaf or island
            section = self._next_root_section()

        return section

    def _previous_parent_section(self):
        # We get it through the "through" class so we can order
        # the parents by weight
        relation = self._parent_relations()[0]
        return relation.parent

    def _previous_root_section(self):
        if self.is_root() or self.is_island():
            root = self
        else:
            root = self.get_roots()[0]

        found = False
        for section in self.sections_in_same_story.order_by('-weight'):
            if found:
                if section.is_root() or section.is_island():
                    return section
            else:
                if section == root:
                    found = True

        return None

    def get_previous_section(self):
        """Get the previous section"""
        # TODO: Unit test for this
        if self.ancestors_set():
            # The section has ancestors.  Return the first one.
            section = self._previous_parent_section()
        else:
            # The section is a root or an island
            section = self._previous_root_section()

        return section

    def render(self, format='html'):
        """Render a representation of the section structure"""
        try:
            return getattr(self, "render_" + format).__call__()
        except AttributeError:
            return self.__unicode__()

    def children_flat(self):
        """
        Return a list of child sections ordered with each branch
        coming before the next
        """
        flattened = []
        for relation in self._child_relations():
            flattened.append(relation.section)
            flattened = flattened + relation.section.children_flat()
        return flattened

    def to_simple(self):
        """
        Return a simplified representation of this object for serialization
        """
        simple = {
            'section_id': self.section_id,
            'title': self.title
        }
        next_section = self.get_next_section()
        previous_section = self.get_previous_section()
        if next_section:
            simple.update({'next_section_id': next_section.section_id})
        if previous_section:
            simple.update(
                {'previous_section_id': previous_section.section_id})

        return simple

    def render_html(self):
        """Render a HTML representation of the section structure"""
        output = []
        output.append('
')
        output.append("
%s
" % self.title)
        if self.assets.count() > 0:
            output.append("
Assets
")
            output.append("
")
            for asset in self.assets.order_by('sectionasset__weight'):
                asset_title = asset.title
                if not asset_title:
                    asset_title = unicode(asset)
                output.append("
%s
" % asset_title)
            output.append("
")

        if self.children.count():
            output.append("
")
            for child_relation in self._child_relations():
                output.append(child_relation.child.render_html())
            output.append("
")

        output.append("
")

        return mark_safe(u'\n'.join(output))

    def change_link(self):
        """Generate a link to the Django admin change page

        You can specify this in the Model Admin's readonly_fields or
        list_display options

        """
        if self.pk:
            change_url = urlresolvers.reverse(
                'admin:storybase_story_section_change', args=(self.pk,))
            return "Change Section" % change_url
        else:
            return ''
    change_link.short_description = 'Change'
    change_link.allow_tags = True

class SectionTranslation(TranslationModel):
    """Translated fields of a Section"""
    section = models.ForeignKey('Section')
    title = ShortTextField()

    class Meta:
        """Model metadata options"""
        unique_together = (('section', 'language'))

    def __unicode__(self):
        return self.title

class SectionRelation(edge_factory(Section, concrete=False)):
    """Through class for parent/child relationships between sections"""
    weight = models.IntegerField(default=0)

class SectionAsset(models.Model):
    """Through class for Asset to Section relations"""
    section = models.ForeignKey('Section')
    asset = models.ForeignKey('storybase_asset.Asset')
    weight = models.IntegerField(default=0)

def add_section_asset_to_story(sender, instance, **kwargs):
    """When an asset is added to a Section, also add it to the Story

    Should be connected to SectionAsset's post_save signal.
    """
    if instance.asset not in instance.section.story.assets.all():
        # An asset was added to a section but it is not related to
        # the section's story.
        # Add it to the Story's list of assets.
        instance.section.story.assets.add(instance.asset)
        instance.section.story.save()

# Add assets to stories when they're added to sections
post_save.connect(add_section_asset_to_story, sender=SectionAsset)

def update_story_last_edited(sender, instance, **kwargs):
    """Update the a section's story's last edited field

    Should be connected to Section's post_save signal.
    """
    # Last edited is automatically set on save
    instance.story.save()

# Update a section's story's last edited field when the section is saved
post_save.connect(update_story_last_edited, sender=Section)

def create_story(title, structure=structure.DEFAULT_STRUCTURE, summary='',
                 language=settings.LANGUAGE_CODE,
                 *args, **kwargs):
    """Convenience function for creating a Story

    Allows for the creation of stories without having to explicitly
    deal with the translations.

    """
    obj = Story(structure=structure, *args, **kwargs)
    obj.save()
    translation = StoryTranslation(story=obj, title=title, summary=summary,
                                   language=language)
    translation.save()
    return obj

def create_section(title, story, language=settings.LANGUAGE_CODE,
                   *args, **kwargs):
    """Convenience function for creating a Section

    Allows for the creation of a section without having to explicitly
    deal with the tranlsations.

    """
    obj = Section(story=story, *args, **kwargs)
    obj.save()
    translation = SectionTranslation(section=obj, title=title,
                                     language=language)
    translation.save()
    return obj
<commit_msg>Fixed another typo.<commit_after>"""Models for stories and story sections"""

from django.conf import settings
from django.contrib.auth.models import User
from django.contrib.sites.models import Site
from django.core import urlresolvers
from django.db import models
from django.db.models.signals import post_save, pre_save
from django.utils import simplejson
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext as _
from django_dag.models import edge_factory, node_factory
# TODO: Decide on tagging suggestion admin app.
# Right now, I'm using a hacked version of
# https://bitbucket.org/fabian/django-taggit-autosuggest
# which I modified to allow for specifying a Tag model
# other than taggit.models.Tag
#from taggit_autosuggest.managers import TaggableManager
from uuidfield.fields import UUIDField
from storybase.fields import ShortTextField
from storybase.models import (LicensedModel, PublishedModel,
    TimestampedModel, TranslatedModel, TranslationModel,
    set_date_on_published)
from storybase.utils import slugify
from storybase_asset.models import Asset
from storybase_user.models import Organization, Project
from storybase_story.managers import StoryManager
from storybase_story import structure
#from storybase_tag.models import TaggedItem

class StoryTranslation(TranslationModel):
    """Encapsulates translated fields of a Story"""
    story = models.ForeignKey('Story')
    title = ShortTextField()
    summary = models.TextField(blank=True)

    class Meta:
        """Model metadata options"""
        unique_together = (('story', 'language'))

    def __unicode__(self):
        return self.title

class Story(TranslatedModel, LicensedModel, PublishedModel,
            TimestampedModel):
    """Metadata for a story

    The Story model stores a story's metadata and aggregates a story's
    media assets

    """
    story_id = UUIDField(auto=True)
    slug = models.SlugField(blank=True)
    byline = models.TextField()
    structure = models.CharField(max_length=20,
        choices=structure.manager.get_structure_options())
    # blank=True, null=True to bypass validation so the user doesn't
    # have to always remember to set this in the Django admin.
    # Though this is set to blank=True, null=True, we should always set
    # this value.  In fact, the StoryModelAdmin class sets this to
    # request.user
    author = models.ForeignKey(User, related_name="stories", blank=True,
                               null=True)
    assets = models.ManyToManyField(Asset, related_name='stories',
                                    blank=True)
    featured_assets = models.ManyToManyField(Asset,
       related_name='featured_in_stories', blank=True,
       help_text=_("Assets to be displayed in teaser version of Story"))
    organizations = models.ManyToManyField(Organization,
                                           related_name='stories',
                                           blank=True)
    projects = models.ManyToManyField(Project, related_name='stories',
                                      blank=True)
    on_homepage = models.BooleanField(_("Featured on homepage"),
		                      default=False)
    #tags = TaggableManager(through=TaggedItem, blank=True)

    objects = StoryManager()

    translated_fields = ['title', 'summary']
    translation_set = 'storytranslation_set'

    _structure_obj = None

    class Meta:
        """Model metadata options"""
        verbose_name_plural = "stories"

    def __unicode__(self):
        return self.title

    @models.permalink
    def get_absolute_url(self):
        """Calculate the canonical URL for a Story"""
        return ('story_detail', [self.slug])

    def get_full_url(self):
        """
        Calculate the canonical URL for a Story, including its
        domain name.
        """
        return "http://%s%s" % (Site.objects.get_current().domain,
                                self.get_absolute_url())

    @property
    def contributor_name(self):
        """
        Return the contributor's first name and last initial or username
        """
        contributor = self.author
        contributor_name = contributor.username
        if contributor.first_name:
            contributor_name = contributor.first_name

            if contributor.last_name:
                contributor_name = "%s %s." % (contributor.first_name,
                                               contributor.last_name[0])
            else:
                contributor_name = contributor.first_name

        return contributor_name

    def get_structure_obj(self):
        """Return a structure object for the story"""
        if self._structure_obj is None:
            # A structure object hasn't been instantiated yet.
            # Create one.
            structure_class = structure.manager.get_structure_class(
            self.structure)
            self._structure_obj = structure_class.__call__(story=self)

        return self._structure_obj

    def sections_flat(self):
        """
        Return a list of sections ordered with each branch coming before
        the next
        """
        flattened = []
        for section in self.sections.filter(root=True).order_by('weight'):
            flattened.append(section)
            flattened = flattened + section.children_flat()

        return flattened

    def sections_json(self):
        """Return a JSON representation of the story sections"""
        sections = []
        for section in self.sections_flat():
            sections.append(section.to_simple())

        return mark_safe(simplejson.dumps(sections))

    def to_simple(self):
        """
        Return a simplified version of the story object that can be
        passed to json.dumps()
        """
        return {
            'story_id': self.story_id,
            'title': self.title,
            'summary': self.summary
        }

    def to_json(self):
        """Return JSON representation of this object"""
        return mark_safe(simplejson.dumps(self.to_simple()))

    def render_featured_asset(self, format='html'):
        """Render a representation of the story's featured asset"""
        try:
            # Return the first featured asset.  We have the ability of
            # selecting multiple featured assets.  Perhaps in the future
            # allow for specifying a particular feature asset or randomly
            # displaying one.
            featured_asset = self.featured_assets.select_subclasses()[0]
            thumbnail_options = {}
            if format == 'html':
                thumbnail_options.update({'html_class': 'featured-asset'})
            return featured_asset.render_thumbnail(format=format,
                                                   **thumbnail_options)
        except IndexError:
            # No featured assets
            return ''

    def featured_asset_thumbnail_url(self):
        """Return the URL of the featured asset's thumbnail

        Returns None if the asset cannot be converted to a thumbnail image.

        """
        try:
            featured_asset = self.featured_assets.select_subclasses()[0]
            return featured_asset.get_thumbnail_url()
        except IndexError:
            # No featured assets
            return None

    def render_story_structure(self, format='html'):
        """Render a representation of the Story structure"""
        output = []
        try:
            root = self.get_root_section()
        except Section.DoesNotExist:
            return ''

        output.append('
')
        output.append(root.render(format))
        output.append('
')
        return mark_safe(u'\n'.join(output))

    def render_toc(self, format='html'):
        """Render a representation of the Story's table of contents"""
        structure = self.get_structure_obj()
        return structure.render_toc(format)

def set_story_slug(sender, instance, **kwargs):
    """
    When a StoryTranslation is saved, set its Story's slug if it doesn't have
    one

    Should be connected to StoryTranslation's post_save signal.
    """
    if not instance.story.slug:
        instance.story.slug = slugify(instance.title)
	instance.story.save()

# Hook up some signal handlers
pre_save.connect(set_date_on_published, sender=Story)
post_save.connect(set_story_slug, sender=StoryTranslation)

class Section(node_factory('SectionRelation'), TranslatedModel):
    """ Section of a story """
    section_id = UUIDField(auto=True)
    story = models.ForeignKey('Story', related_name='sections')
    # True if this section the root section of the story, either
    # the first section in a linear story, or the central node
    # in a drill-down/"spider" structure.  Otherwise, False
    root = models.BooleanField(default=False)
    weight = models.IntegerField(default=0)
    """The ordering of top-level sections relative to each other"""
    assets = models.ManyToManyField(Asset, related_name='sections',
                                    blank=True, through='SectionAsset')

    translated_fields = ['title']
    translation_set = 'sectiontranslation_set'

    def __unicode__(self):
        try:
            return self.title
        except IndexError:
            # HACK: Need this to support deleting Sections through
            # an inline on the Story Change page
            #
            # When deleting an object in the Django admin, a Section
            # object gets instantiated with no translation set.
            # So, the call to __getattr__() raises an IndexError
            return super(Section, self).__unicode__()

    @property
    def sections_in_same_story(self):
        """Return other sections in the same story"""
        return self.__class__.objects.filter(story=self.story)

    def _child_relations(self):
        """
        Get a query set of through model instances for children
        of this section
        """
        return self.children.through.objects.filter(parent=self).order_by(
            'weight', 'child__sectiontranslation__title')

    def _parent_relations(self):
        """
        Get a list of through model instances for parents of this
        section
        """
        # This is a little confusing because we want parents, not
        # children.  self.children.through just gets us the "through" class
        # for the relation.
        return self.children.through.objects.filter(child=self).order_by(
            'weight', 'child__sectiontranslation__title')

    # TODO: Think about the best strategy for building a sense of
    # previous and next sections. I'm not sure if this is the best
    # place for the functionality and there's likely a way to avoid
    # having to traverse the tree so many times.

    def _next_child_section(self):
        """Return the first section that is a child of this section"""
        # We get it through the "through" class so we can order
        # the children by weight
        relation = self._child_relations()[0]
        return relation.child

    def _next_root_section(self):
        """Return the next top-level section"""
        if self.is_root() or self.is_island():
            root = self
        else:
            root = self.get_roots()[0]

        found = False
        for section in self.sections_in_same_story.order_by('weight'):
            if found:
                if section.is_root() or section.is_island():
                    return section
            else:
                if section == root:
                    found = True

        return None

    def get_next_section(self):
        """Get the next section"""
        # TODO: Unit test for this
        if self.children.count():
            # The section has children, return the first child
            section = self._next_child_section()
        else:
            # The section is a leaf or island
            section = self._next_root_section()

        return section

    def _previous_parent_section(self):
        # We get it through the "through" class so we can order
        # the parents by weight
        relation = self._parent_relations()[0]
        return relation.parent

    def _previous_root_section(self):
        if self.is_root() or self.is_island():
            root = self
        else:
            root = self.get_roots()[0]

        found = False
        for section in self.sections_in_same_story.order_by('-weight'):
            if found:
                if section.is_root() or section.is_island():
                    return section
            else:
                if section == root:
                    found = True

        return None

    def get_previous_section(self):
        """Get the previous section"""
        # TODO: Unit test for this
        if self.ancestors_set():
            # The section has ancestors.  Return the first one.
            section = self._previous_parent_section()
        else:
            # The section is a root or an island
            section = self._previous_root_section()

        return section

    def render(self, format='html'):
        """Render a representation of the section structure"""
        try:
            return getattr(self, "render_" + format).__call__()
        except AttributeError:
            return self.__unicode__()

    def children_flat(self):
        """
        Return a list of child sections ordered with each branch
        coming before the next
        """
        flattened = []
        for relation in self._child_relations():
            flattened.append(relation.child)
            flattened = flattened + relation.child.children_flat()
        return flattened

    def to_simple(self):
        """
        Return a simplified representation of this object for serialization
        """
        simple = {
            'section_id': self.section_id,
            'title': self.title
        }
        next_section = self.get_next_section()
        previous_section = self.get_previous_section()
        if next_section:
            simple.update({'next_section_id': next_section.section_id})
        if previous_section:
            simple.update(
                {'previous_section_id': previous_section.section_id})

        return simple

    def render_html(self):
        """Render a HTML representation of the section structure"""
        output = []
        output.append('
')
        output.append("
%s
" % self.title)
        if self.assets.count() > 0:
            output.append("
Assets
")
            output.append("
")
            for asset in self.assets.order_by('sectionasset__weight'):
                asset_title = asset.title
                if not asset_title:
                    asset_title = unicode(asset)
                output.append("
%s
" % asset_title)
            output.append("
")

        if self.children.count():
            output.append("
")
            for child_relation in self._child_relations():
                output.append(child_relation.child.render_html())
            output.append("
")

        output.append("
")

        return mark_safe(u'\n'.join(output))

    def change_link(self):
        """Generate a link to the Django admin change page

        You can specify this in the Model Admin's readonly_fields or
        list_display options

        """
        if self.pk:
            change_url = urlresolvers.reverse(
                'admin:storybase_story_section_change', args=(self.pk,))
            return "Change Section" % change_url
        else:
            return ''
    change_link.short_description = 'Change'
    change_link.allow_tags = True

class SectionTranslation(TranslationModel):
    """Translated fields of a Section"""
    section = models.ForeignKey('Section')
    title = ShortTextField()

    class Meta:
        """Model metadata options"""
        unique_together = (('section', 'language'))

    def __unicode__(self):
        return self.title

class SectionRelation(edge_factory(Section, concrete=False)):
    """Through class for parent/child relationships between sections"""
    weight = models.IntegerField(default=0)

class SectionAsset(models.Model):
    """Through class for Asset to Section relations"""
    section = models.ForeignKey('Section')
    asset = models.ForeignKey('storybase_asset.Asset')
    weight = models.IntegerField(default=0)

def add_section_asset_to_story(sender, instance, **kwargs):
    """When an asset is added to a Section, also add it to the Story

    Should be connected to SectionAsset's post_save signal.
    """
    if instance.asset not in instance.section.story.assets.all():
        # An asset was added to a section but it is not related to
        # the section's story.
        # Add it to the Story's list of assets.
        instance.section.story.assets.add(instance.asset)
        instance.section.story.save()

# Add assets to stories when they're added to sections
post_save.connect(add_section_asset_to_story, sender=SectionAsset)

def update_story_last_edited(sender, instance, **kwargs):
    """Update the a section's story's last edited field

    Should be connected to Section's post_save signal.
    """
    # Last edited is automatically set on save
    instance.story.save()

# Update a section's story's last edited field when the section is saved
post_save.connect(update_story_last_edited, sender=Section)

def create_story(title, structure=structure.DEFAULT_STRUCTURE, summary='',
                 language=settings.LANGUAGE_CODE,
                 *args, **kwargs):
    """Convenience function for creating a Story

    Allows for the creation of stories without having to explicitly
    deal with the translations.

    """
    obj = Story(structure=structure, *args, **kwargs)
    obj.save()
    translation = StoryTranslation(story=obj, title=title, summary=summary,
                                   language=language)
    translation.save()
    return obj

def create_section(title, story, language=settings.LANGUAGE_CODE,
                   *args, **kwargs):
    """Convenience function for creating a Section

    Allows for the creation of a section without having to explicitly
    deal with the tranlsations.

    """
    obj = Section(story=story, *args, **kwargs)
    obj.save()
    translation = SectionTranslation(section=obj, title=title,
                                     language=language)
    translation.save()
    return obj
<|endoftext|>