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

import sys
import os.path

from datetime import datetime

import re
import json
import csv

from lxml import etree
from dateutil import tz
import dateutil.parser
from copy import deepcopy

from . import crawler

from wheezy.template.engine import Engine
from wheezy.template.loader import DictLoader
from wheezy.template.ext.core import CoreExtension

json.encoder.c_make_encoder = None

try:
    from collections import OrderedDict
except ImportError:
    # python < 2.7
    from ordereddict import OrderedDict

try:
    from StringIO import StringIO
    from urllib2 import urlopen
    from ConfigParser import ConfigParser
except ImportError:
    # python > 3
    from io import StringIO
    from urllib.request import urlopen
    from configparser import ConfigParser

try:
    basestring
except NameError:
    basestring = unicode = str

Element = etree.Element

NSMAP = {'atom': 'http://www.w3.org/2005/Atom',
         'atom03': 'http://purl.org/atom/ns#',
         'media': 'http://search.yahoo.com/mrss/',
         'rdf': 'http://www.w3.org/1999/02/22-rdf-syntax-ns#',
         'slash': 'http://purl.org/rss/1.0/modules/slash/',
         'dc': 'http://purl.org/dc/elements/1.1/',
         'content': 'http://purl.org/rss/1.0/modules/content/',
         'rssfake': 'http://purl.org/rss/1.0/'}

def load(url):
    d = urlopen(url).read()
    return parse(d)

def tag_NS(tag, nsmap=NSMAP):
    match = re.search(r'^\{([^\}]+)\}(.*)$', tag)
    if match:
        match = match.groups()
        for (key, url) in nsmap.items():
            if url == match[0]:
                return "%s:%s" % (key, match[1].lower())
    else:
        match = re.search(r'^([^:]+):([^:]+)$', tag)
        if match:
            match = match.groups()
            if match[0] in nsmap:
                return "{%s}%s" % (nsmap[match[0]], match[1].lower())
    return tag

def inner_html(xml):
    return (xml.text or '') + ''.join([etree.tostring(child) for child in xml])

def clean_node(xml):
    [xml.remove(child) for child in xml]

def parse_rules(filename=None):
    if not filename:
        filename = os.path.join(os.path.dirname(__file__), 'feedify.ini')

    config = ConfigParser()
    config.read(filename)

    rules = dict([(x, dict(config.items(x))) for x in config.sections()])

    for section in rules.keys():
        for arg in rules[section].keys():
            if '\n' in rules[section][arg]:
                rules[section][arg] = rules[section][arg].split('\n')[1:]

    return rules

class ParserBase(object):
    def __init__(self, data=None, rules=None):
        if rules is None:
            rules = parse_rules()['rss']

        if data is None:
            data = rules['base'][0]

        self.rules = rules
        self.root = self.parse(data)

        # do `if multi` and select the correct rule for each (and split \n)
        if isinstance(self.rules['items'], list):
            for (i, rule) in enumerate(self.rules['items']):
                if self.rule_search(rule) is not None:
                    key = i
                    break

            else:
                key = 0

            for arg in self.rules.keys():
                if isinstance(self.rules[arg], list):
                    self.rules[arg] = self.rules[arg][key]

    def parse(self, raw):
        pass

    def tostring(self):
        # output in its input format
        # to output in sth fancy (json, csv, html), change class type
        pass

    def tojson(self, indent=None):
        # TODO temporary
        return json.dumps(OrderedDict(self.iterdic()), indent=indent)

    def tocsv(self):
        # TODO temporary
        out = StringIO()
        c = csv.writer(out, dialect=csv.excel)

        for item in self.items:
            row = [getattr(item, x) for x in item.dic]

            if sys.version_info[0] < 3:
                row = [x.encode('utf-8') if isinstance(x, unicode) else x for x in row]

            c.writerow(row)

        out.seek(0)
        return out.read()

    def tohtml(self):
        # TODO temporary
        path = os.path.join(os.path.dirname(__file__), 'reader.html.template')
        loader = DictLoader({'reader': open(path).read()})
        engine = Engine(loader=loader, extensions=[CoreExtension()])
        template = engine.get_template('reader')
        return template.render({'feed': self}).encode('utf-8')

    def iterdic(self):
        for element in self.dic:
            value = getattr(self, element)

            if element == 'items':
                value = [OrderedDict(x.iterdic()) for x in value]
            elif isinstance(value, datetime):
                value = value.isoformat()

            yield element, value

    def rule_search(self, rule):
        # xpath, return the first one only
        try:
            return self.rule_search_all(rule)[0]

        except IndexError:
            return None

    def rule_search_all(self, rule):
        # xpath, return all (useful to find feed items)
        pass

    def rule_search_last(self, rule):
        # xpath, return the first one only
        try:
            return self.rule_search_all(rule)[-1]

        except IndexError:
            return None

    def rule_create(self, rule):
        # create node based on rule
        # (duplicate, copy existing (or template) or create from scratch, if possible)
        # --> might want to create node_duplicate helper fns
        pass

    def rule_remove(self, rule):
        # remove node from its parent
        pass

    def rule_set(self, rule, value):
        # value is always a str?
        pass

    def rule_str(self, rule):
        # GETs inside (pure) text from it
        pass

    def bool_prs(self, x):
        # parse
        pass

    def bool_fmt(self, x):
        # format
        pass

    def time_prs(self, x):
        # parse
        pass

    def time_fmt(self, x):
        # format
        pass

    def get_raw(self, rule_name):
        # get the raw output, for self.get_raw('items')
        pass

    def get_str(self, rule_name):
        # simple function to get nice text from the rule name
        # for use in @property, ie. self.get_str('title')
        pass

    def set_str(self, rule_name):
        pass

    def remove(self, rule_name):
        # easy deleter
        pass

class ParserXML(ParserBase):
    def parse(self, raw):
        parser = etree.XMLParser(recover=True)
        return etree.fromstring(raw, parser)

    def tostring(self, **k):
        return etree.tostring(self.root, **k)

    def _rule_parse(self, rule):
        test = re.search(r'^(.*)/@([a-z]+)$', rule) # to match //div/a/@href
        return test.groups() if test else (rule, None)

    def _resolve_ns(self, rule):
        match = re.search(r'^([^:]+):([^:]+)$', rule)
        if match:
            match = match.groups()
            if match[0] in NSMAP:
                return "{%s}%s" % (NSMAP[match[0]], match[1].lower())

        return rule

    def rule_search_all(self, rule):
        try:
            return self.root.xpath(rule, namespaces=NSMAP)

        except etree.XPathEvalError:
            return []

    def rule_create(self, rule):
        # duplicate, copy from template or create from scratch
        rule, key = self._rule_parse(rule)

        # try recreating based on the rule (for really basic rules, ie. plain RSS)
        if re.search(r'^[a-zA-Z0-9/:]+$', rule):
            chain = rule.strip('/').split('/')
            current = self.root

            if rule[0] == '/':
                chain = chain[1:]

            for (i, node) in enumerate(chain):
                test = current.find(self._resolve_ns(node))

                if test and i < len(chain) - 1:
                    # yay, go on
                    current = test

                else:
                    # opps need to create
                    element = etree.Element(self._resolve_ns(node))
                    current.append(element)
                    current = element

            return current

        # try duplicating from existing (works well with fucked up structures)
        match = self.rule_search_last(rule)
        if match:
            element = deepcopy(match)
            match.getparen().append(element)
            return element

        # try duplicating from template
        # FIXME
        # >>> self.xml.getroottree().getpath(ff.find('a'))

        return None

    def rule_remove(self, rule):
        rule, key = self._rule_parse(rule)

        match = self.rule_search(rule)

        if key is not None:
            del x.attrib[key]

        else:
            match.getparent().remove(match)

    def rule_set(self, rule, value):
        rule, key = self._rule_parse(rule)

        match = self.rule_search(rule)

        if key is not None:
            match.attrib[key] = value

        else:
            match.text = value

    def rule_str(self, rule):
        match = self.rule_search(rule)

        if isinstance(match, etree._Element):
            return match.text or ""

        else:
            return match or ""

    def bool_prs(self, x):
        return (x or '').lower() != 'false'

    def bool_fmt(self, x):
        return 'true' if x else 'false'

    def time_prs(self, x):
        try:
            return parse_time(x)
        except ValueError:
            return None

    def time_fmt(self, x):
        try:
            time = parse_time(x)
            return time.strftime(self.rules['timeformat'])
        except ValueError:
            pass

    def get_raw(self, rule_name):
        return self.rule_search_all(self.rules[rule_name])

    def get_str(self, rule_name):
        return self.rule_str(self.rules[rule_name])

    def set_str(self, rule_name, value):
        try:
            return self.rule_set(self.rules[rule_name], value)

        except AttributeError:
            # does not exist, have to create it
            self.rule_create(self.rules[rule_name])
            return self.rule_set(self.rules[rule_name], value)

    def remove(self, rule_name):
        self.rule_remove(self.rules[rule_name])

def parse_time(value):
    if isinstance(value, basestring):
        if re.match(r'^[0-9]+$', value):
            return datetime.fromtimestamp(int(value), tz.tzutc())
        else:
            return dateutil.parser.parse(value, tzinfos=tz.tzutc)
    elif isinstance(value, int):
        return datetime.fromtimestamp(value, tz.tzutc())
    elif isinstance(value, datetime):
        return value
    else:
        return False

class Uniq(object):
    _map = {}
    _id = None

    def __new__(cls, *args, **kwargs):
        # check if an item was already created for it
        # if so, reuse it
        # if not, create a new one

        tmp_id = cls._gen_id(*args, **kwargs)
        if tmp_id is not None and tmp_id in cls._map:
            return cls._map[tmp_id]

        else:
            obj = object.__new__(cls, *args, **kwargs)
            cls._map[obj._id] = obj
            return obj

class Feed(object):
    itemsClass = 'Item'
    dic = ('title', 'desc', 'items')

    def wrap_items(self, items):
        itemsClass = globals()[self.itemsClass]
        return [itemsClass(x, self.rules) for x in items]

    title = property(
        lambda f:   f.get_str('title'),
        lambda f,x: f.set_str('title', x),
        lambda f:   f.remove('title') )
    description = desc = property(
        lambda f:   f.get_str('desc'),
        lambda f,x: f.set_str('desc', x),
        lambda f:   f.remove('desc') )
    items = property(
        lambda f:   f )

    def append(self, new=None):
        self.rule_create(self.rules['items'])
        item = self.items[-1]

        if new is None:
            return item

        for attr in globals()[self.itemsClass].dic:
            if hasattr(new, attr):
                setattr(element, attr, getattr(new, attr))

            elif attr in cousin:
                setattr(element, attr, new[attr])

    def __getitem__(self, key):
        return self.wrap_items(self.get_raw('items'))[key]

    def __delitem__(self, key):
        self[key].remove()

    def __len__(self):
        return len(self.get_raw('items'))

class FeedXML(Feed, ParserXML):
    itemsClass = 'ItemXML'

    def tostring(self, **k):
        return etree.tostring(self.root.getroottree(), **k)

class Item(Uniq):
    dic = ('title', 'link', 'desc', 'content', 'id', 'is_permalink', 'time', 'updated')

    def __init__(self, xml=None, rules=None):
        self._id = self._gen_id(xml)
        self.root = xml
        self.rules = rules

    @staticmethod
    def _gen_id(xml=None, *args, **kwargs):
        return id(xml)

    title = property(
        lambda f:   f.get_str('item_title'),
        lambda f,x: f.set_str('item_title', x),
        lambda f:   f.remove('item_title') )
    link = property(
        lambda f:   f.get_str('item_link'),
        lambda f,x: f.set_str('item_link', x),
        lambda f:   f.remove('item_link') )
    description = desc = property(
        lambda f:   f.get_str('item_desc'),
        lambda f,x: f.set_str('item_desc', x),
        lambda f:   f.remove('item_desc') )
    content = property(
        lambda f:   f.get_str('item_content'),
        lambda f,x: f.set_str('item_content', x),
        lambda f:   f.remove('item_content') )
    id = property(
        lambda f:   f.get_str('item_id'),
        lambda f,x: f.set_str('item_id', x),
        lambda f:   f.remove('item_id') )
    is_permalink = property(
        lambda f:   f.get_str('item_is_permalink'),
        lambda f,x: f.set_str('item_is_permalink', x))#,
        #lambda f:   f.remove('item_is_permalink') )
    time = property(
        lambda f:   f.time_fmt(f.get_str('item_time')),
        lambda f,x: f.set_str('title', f.time_prs(x)),
        lambda f:   f.remove('item_time') )
    updated = property(
        lambda f:   f.time_fmt(f.get_str('item_updated')),
        lambda f,x: f.set_str('updated', f.time_prs(x)),
        lambda f:   f.remove('item_updated') )

class ItemXML(Item, ParserXML):
    pass
<commit_msg>feeds: fix atom xhtml handling<commit_after>#!/usr/bin/env python

import sys
import os.path

from datetime import datetime

import re
import json
import csv

from lxml import etree
from dateutil import tz
import dateutil.parser
from copy import deepcopy

from . import crawler

from wheezy.template.engine import Engine
from wheezy.template.loader import DictLoader
from wheezy.template.ext.core import CoreExtension

json.encoder.c_make_encoder = None

try:
    from collections import OrderedDict
except ImportError:
    # python < 2.7
    from ordereddict import OrderedDict

try:
    from StringIO import StringIO
    from urllib2 import urlopen
    from ConfigParser import ConfigParser
except ImportError:
    # python > 3
    from io import StringIO
    from urllib.request import urlopen
    from configparser import ConfigParser

try:
    basestring
except NameError:
    basestring = unicode = str

Element = etree.Element

NSMAP = {'atom': 'http://www.w3.org/2005/Atom',
         'atom03': 'http://purl.org/atom/ns#',
         'media': 'http://search.yahoo.com/mrss/',
         'rdf': 'http://www.w3.org/1999/02/22-rdf-syntax-ns#',
         'slash': 'http://purl.org/rss/1.0/modules/slash/',
         'dc': 'http://purl.org/dc/elements/1.1/',
         'content': 'http://purl.org/rss/1.0/modules/content/',
         'rssfake': 'http://purl.org/rss/1.0/'}

def load(url):
    d = urlopen(url).read()
    return parse(d)

def tag_NS(tag, nsmap=NSMAP):
    match = re.search(r'^\{([^\}]+)\}(.*)$', tag)
    if match:
        match = match.groups()
        for (key, url) in nsmap.items():
            if url == match[0]:
                return "%s:%s" % (key, match[1].lower())
    else:
        match = re.search(r'^([^:]+):([^:]+)$', tag)
        if match:
            match = match.groups()
            if match[0] in nsmap:
                return "{%s}%s" % (nsmap[match[0]], match[1].lower())
    return tag

def parse_rules(filename=None):
    if not filename:
        filename = os.path.join(os.path.dirname(__file__), 'feedify.ini')

    config = ConfigParser()
    config.read(filename)

    rules = dict([(x, dict(config.items(x))) for x in config.sections()])

    for section in rules.keys():
        for arg in rules[section].keys():
            if '\n' in rules[section][arg]:
                rules[section][arg] = rules[section][arg].split('\n')[1:]

    return rules

class ParserBase(object):
    def __init__(self, data=None, rules=None):
        if rules is None:
            rules = parse_rules()['rss']

        if data is None:
            data = rules['base'][0]

        self.rules = rules
        self.root = self.parse(data)

        # do `if multi` and select the correct rule for each (and split \n)
        if isinstance(self.rules['items'], list):
            for (i, rule) in enumerate(self.rules['items']):
                if self.rule_search(rule) is not None:
                    key = i
                    break

            else:
                key = 0

            for arg in self.rules.keys():
                if isinstance(self.rules[arg], list):
                    self.rules[arg] = self.rules[arg][key]

    def parse(self, raw):
        pass

    def tostring(self):
        # output in its input format
        # to output in sth fancy (json, csv, html), change class type
        pass

    def tojson(self, indent=None):
        # TODO temporary
        return json.dumps(OrderedDict(self.iterdic()), indent=indent)

    def tocsv(self):
        # TODO temporary
        out = StringIO()
        c = csv.writer(out, dialect=csv.excel)

        for item in self.items:
            row = [getattr(item, x) for x in item.dic]

            if sys.version_info[0] < 3:
                row = [x.encode('utf-8') if isinstance(x, unicode) else x for x in row]

            c.writerow(row)

        out.seek(0)
        return out.read()

    def tohtml(self):
        # TODO temporary
        path = os.path.join(os.path.dirname(__file__), 'reader.html.template')
        loader = DictLoader({'reader': open(path).read()})
        engine = Engine(loader=loader, extensions=[CoreExtension()])
        template = engine.get_template('reader')
        return template.render({'feed': self}).encode('utf-8')

    def iterdic(self):
        for element in self.dic:
            value = getattr(self, element)

            if element == 'items':
                value = [OrderedDict(x.iterdic()) for x in value]
            elif isinstance(value, datetime):
                value = value.isoformat()

            yield element, value

    def rule_search(self, rule):
        # xpath, return the first one only
        try:
            return self.rule_search_all(rule)[0]

        except IndexError:
            return None

    def rule_search_all(self, rule):
        # xpath, return all (useful to find feed items)
        pass

    def rule_search_last(self, rule):
        # xpath, return the first one only
        try:
            return self.rule_search_all(rule)[-1]

        except IndexError:
            return None

    def rule_create(self, rule):
        # create node based on rule
        # (duplicate, copy existing (or template) or create from scratch, if possible)
        # --> might want to create node_duplicate helper fns
        pass

    def rule_remove(self, rule):
        # remove node from its parent
        pass

    def rule_set(self, rule, value):
        # value is always a str?
        pass

    def rule_str(self, rule):
        # GETs inside (pure) text from it
        pass

    def bool_prs(self, x):
        # parse
        pass

    def bool_fmt(self, x):
        # format
        pass

    def time_prs(self, x):
        # parse
        pass

    def time_fmt(self, x):
        # format
        pass

    def get_raw(self, rule_name):
        # get the raw output, for self.get_raw('items')
        pass

    def get_str(self, rule_name):
        # simple function to get nice text from the rule name
        # for use in @property, ie. self.get_str('title')
        pass

    def set_str(self, rule_name):
        pass

    def remove(self, rule_name):
        # easy deleter
        pass

class ParserXML(ParserBase):
    def parse(self, raw):
        parser = etree.XMLParser(recover=True)
        return etree.fromstring(raw, parser)

    def tostring(self, **k):
        return etree.tostring(self.root, **k)

    def _rule_parse(self, rule):
        test = re.search(r'^(.*)/@([a-z]+)$', rule) # to match //div/a/@href
        return test.groups() if test else (rule, None)

    def _resolve_ns(self, rule):
        match = re.search(r'^([^:]+):([^:]+)$', rule)
        if match:
            match = match.groups()
            if match[0] in NSMAP:
                return "{%s}%s" % (NSMAP[match[0]], match[1].lower())

        return rule

    @staticmethod
    def _inner_html(xml):
        return (xml.text or '') + ''.join([etree.tostring(child) for child in xml])

    @staticmethod
    def _clean_node(xml):
        [xml.remove(child) for child in xml]

    def rule_search_all(self, rule):
        try:
            return self.root.xpath(rule, namespaces=NSMAP)

        except etree.XPathEvalError:
            return []

    def rule_create(self, rule):
        # duplicate, copy from template or create from scratch
        rule, key = self._rule_parse(rule)

        # try recreating based on the rule (for really basic rules, ie. plain RSS)
        if re.search(r'^[a-zA-Z0-9/:]+$', rule):
            chain = rule.strip('/').split('/')
            current = self.root

            if rule[0] == '/':
                chain = chain[1:]

            for (i, node) in enumerate(chain):
                test = current.find(self._resolve_ns(node))

                if test and i < len(chain) - 1:
                    # yay, go on
                    current = test

                else:
                    # opps need to create
                    element = etree.Element(self._resolve_ns(node))
                    current.append(element)
                    current = element

            return current

        # try duplicating from existing (works well with fucked up structures)
        match = self.rule_search_last(rule)
        if match:
            element = deepcopy(match)
            match.getparen().append(element)
            return element

        # try duplicating from template
        # FIXME
        # >>> self.xml.getroottree().getpath(ff.find('a'))

        return None

    def rule_remove(self, rule):
        rule, key = self._rule_parse(rule)

        match = self.rule_search(rule)

        if key is not None:
            del x.attrib[key]

        else:
            match.getparent().remove(match)

    def rule_set(self, rule, value):
        rule, key = self._rule_parse(rule)

        match = self.rule_search(rule)

        if key is not None:
            match.attrib[key] = value

        else:
            if len(match):
                # atom stuff
                self._clean_node(match)

                if match.attrib.get('type', '') == 'xhtml':
                    match.attrib['type'] = 'html'

            match.text = value

    def rule_str(self, rule):
        match = self.rule_search(rule)

        if isinstance(match, etree._Element):
            if len(match):
                # atom stuff
                return self._inner_html(match)

            else:
                return match.text or ""

        else:
            return match or ""

    def bool_prs(self, x):
        return (x or '').lower() != 'false'

    def bool_fmt(self, x):
        return 'true' if x else 'false'

    def time_prs(self, x):
        try:
            return parse_time(x)
        except ValueError:
            return None

    def time_fmt(self, x):
        try:
            time = parse_time(x)
            return time.strftime(self.rules['timeformat'])
        except ValueError:
            pass

    def get_raw(self, rule_name):
        return self.rule_search_all(self.rules[rule_name])

    def get_str(self, rule_name):
        return self.rule_str(self.rules[rule_name])

    def set_str(self, rule_name, value):
        try:
            return self.rule_set(self.rules[rule_name], value)

        except AttributeError:
            # does not exist, have to create it
            self.rule_create(self.rules[rule_name])
            return self.rule_set(self.rules[rule_name], value)

    def remove(self, rule_name):
        self.rule_remove(self.rules[rule_name])

def parse_time(value):
    if isinstance(value, basestring):
        if re.match(r'^[0-9]+$', value):
            return datetime.fromtimestamp(int(value), tz.tzutc())
        else:
            return dateutil.parser.parse(value, tzinfos=tz.tzutc)
    elif isinstance(value, int):
        return datetime.fromtimestamp(value, tz.tzutc())
    elif isinstance(value, datetime):
        return value
    else:
        return False

class Uniq(object):
    _map = {}
    _id = None

    def __new__(cls, *args, **kwargs):
        # check if an item was already created for it
        # if so, reuse it
        # if not, create a new one

        tmp_id = cls._gen_id(*args, **kwargs)
        if tmp_id is not None and tmp_id in cls._map:
            return cls._map[tmp_id]

        else:
            obj = object.__new__(cls, *args, **kwargs)
            cls._map[obj._id] = obj
            return obj

class Feed(object):
    itemsClass = 'Item'
    dic = ('title', 'desc', 'items')

    def wrap_items(self, items):
        itemsClass = globals()[self.itemsClass]
        return [itemsClass(x, self.rules) for x in items]

    title = property(
        lambda f:   f.get_str('title'),
        lambda f,x: f.set_str('title', x),
        lambda f:   f.remove('title') )
    description = desc = property(
        lambda f:   f.get_str('desc'),
        lambda f,x: f.set_str('desc', x),
        lambda f:   f.remove('desc') )
    items = property(
        lambda f:   f )

    def append(self, new=None):
        self.rule_create(self.rules['items'])
        item = self.items[-1]

        if new is None:
            return item

        for attr in globals()[self.itemsClass].dic:
            if hasattr(new, attr):
                setattr(element, attr, getattr(new, attr))

            elif attr in cousin:
                setattr(element, attr, new[attr])

    def __getitem__(self, key):
        return self.wrap_items(self.get_raw('items'))[key]

    def __delitem__(self, key):
        self[key].remove()

    def __len__(self):
        return len(self.get_raw('items'))

class FeedXML(Feed, ParserXML):
    itemsClass = 'ItemXML'

    def tostring(self, **k):
        return etree.tostring(self.root.getroottree(), **k)

class Item(Uniq):
    dic = ('title', 'link', 'desc', 'content', 'id', 'is_permalink', 'time', 'updated')

    def __init__(self, xml=None, rules=None):
        self._id = self._gen_id(xml)
        self.root = xml
        self.rules = rules

    @staticmethod
    def _gen_id(xml=None, *args, **kwargs):
        return id(xml)

    title = property(
        lambda f:   f.get_str('item_title'),
        lambda f,x: f.set_str('item_title', x),
        lambda f:   f.remove('item_title') )
    link = property(
        lambda f:   f.get_str('item_link'),
        lambda f,x: f.set_str('item_link', x),
        lambda f:   f.remove('item_link') )
    description = desc = property(
        lambda f:   f.get_str('item_desc'),
        lambda f,x: f.set_str('item_desc', x),
        lambda f:   f.remove('item_desc') )
    content = property(
        lambda f:   f.get_str('item_content'),
        lambda f,x: f.set_str('item_content', x),
        lambda f:   f.remove('item_content') )
    id = property(
        lambda f:   f.get_str('item_id'),
        lambda f,x: f.set_str('item_id', x),
        lambda f:   f.remove('item_id') )
    is_permalink = property(
        lambda f:   f.get_str('item_is_permalink'),
        lambda f,x: f.set_str('item_is_permalink', x))#,
        #lambda f:   f.remove('item_is_permalink') )
    time = property(
        lambda f:   f.time_fmt(f.get_str('item_time')),
        lambda f,x: f.set_str('title', f.time_prs(x)),
        lambda f:   f.remove('item_time') )
    updated = property(
        lambda f:   f.time_fmt(f.get_str('item_updated')),
        lambda f,x: f.set_str('updated', f.time_prs(x)),
        lambda f:   f.remove('item_updated') )

class ItemXML(Item, ParserXML):
    pass
<|endoftext|>


<commit_before>from numpy.random.mtrand import uniform
from netCDF4 import Dataset
from numpy.testing import assert_almost_equal
import os, tempfile, unittest

ndim = 100000
filename = tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name
datarr = uniform(size=(ndim,))

def write_netcdf(filename,dtype='f8'):
    nc = Dataset(filename,'w')
    nc.createDimension('n', ndim)
    foo = nc.createVariable('data',\
            dtype,('n'),compression=None)
    foo_szip = nc.createVariable('data_szip',\
            dtype,('n'),compression='szip',szip_coding='ec',szip_pixels_per_block=32)
    foo[:] = datarr
    foo_szip[:] = datarr
    nc.close()

class CompressionTestCase(unittest.TestCase):

    def setUp(self):
        self.filename = filename
        write_netcdf(self.filename)

    def tearDown(self):
        # Remove the temporary files
        os.remove(self.filename)

    def runTest(self):
        f = Dataset(self.filename)
        assert_almost_equal(datarr,f.variables['data'][:])
        assert f.variables['data'].filters() ==\
        {'zlib':False,'szip':False,'zstd':False,'bzip2':False,'blosc':False,'shuffle':False,'complevel':0,'fletcher32':False}
        assert_almost_equal(datarr,f.variables['data_szip'][:])
        dtest = {'zlib': False, 'szip':
                {'coding': 'ec', 'pixels_per_block': 32},
                'zstd': False, 'bzip2': False, 'blosc': False,
                'shuffle': False, 'complevel': 4, 'fletcher32': False}
        print(f.variables['data_szip'].filters())
        assert f.variables['data_szip'].filters() == dtest
        f.close()

if __name__ == '__main__':
    unittest.main()
<commit_msg>update<commit_after>from numpy.random.mtrand import uniform
from netCDF4 import Dataset
from numpy.testing import assert_almost_equal
import os, tempfile, unittest

ndim = 100000
filename = tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name
datarr = uniform(size=(ndim,))

def write_netcdf(filename,dtype='f8'):
    nc = Dataset(filename,'w')
    nc.createDimension('n', ndim)
    foo = nc.createVariable('data',\
            dtype,('n'),compression=None)
    foo_szip = nc.createVariable('data_szip',\
            dtype,('n'),compression='szip',szip_coding='ec',szip_pixels_per_block=32)
    foo[:] = datarr
    foo_szip[:] = datarr
    nc.close()

class CompressionTestCase(unittest.TestCase):

    def setUp(self):
        self.filename = filename
        write_netcdf(self.filename)

    def tearDown(self):
        # Remove the temporary files
        os.remove(self.filename)

    def runTest(self):
        f = Dataset(self.filename)
        assert_almost_equal(datarr,f.variables['data'][:])
        assert f.variables['data'].filters() ==\
        {'zlib':False,'szip':False,'zstd':False,'bzip2':False,'blosc':False,'shuffle':False,'complevel':0,'fletcher32':False}
        assert_almost_equal(datarr,f.variables['data_szip'][:])
        dtest = {'zlib': False, 'szip':
                {'coding': 'ec', 'pixels_per_block': 32},
                'zstd': False, 'bzip2': False, 'blosc': False,
                'shuffle': False, 'complevel': 4, 'fletcher32': False}
        print(f.variables['data_szip'].filters())
        #assert f.variables['data_szip'].filters() == dtest
        f.close()

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>"""
Handles URL downloads.

This is the low-level interface for downloading interfaces, implementations, icons, etc.

@see: L{fetch} higher-level API for downloads that uses this module
"""

# Copyright (C) 2006, Thomas Leonard
# See the README file for details, or visit http://0install.net.

import tempfile, os, sys, subprocess
from zeroinstall import SafeException
from zeroinstall.support import tasks
from logging import info

download_starting = "starting"	# Waiting for UI to start it
download_fetching = "fetching"	# In progress
download_complete = "complete"	# Downloaded and cached OK
download_failed = "failed"

class DownloadError(SafeException):
	"""Download process failed."""
	pass

class DownloadAborted(DownloadError):
	"""Download aborted because of a call to L{Download.abort}"""
	def __init__(self, message):
		SafeException.__init__(self, message or "Download aborted at user's request")

class Download(object):
	"""A download of a single resource to a temporary file.
	@ivar url: the URL of the resource being fetched
	@type url: str
	@ivar tempfile: the file storing the downloaded data
	@type tempfile: file
	@ivar status: the status of the download
	@type status: (download_starting | download_fetching | download_failed | download_complete)
	@ivar errors: data received from the child's stderr
	@type errors: str
	@ivar expected_size: the expected final size of the file
	@type expected_size: int | None
	@ivar downloaded: triggered when the download ends (on success or failure)
	@type downloaded: L{tasks.Blocker}
	@ivar hint: hint passed by and for caller
	@type hint: object
	@ivar child: the child process
	@type child: L{subprocess.Popen}
	@ivar aborted_by_user: whether anyone has called L{abort}
	@type aborted_by_user: bool
	"""
	__slots__ = ['url', 'tempfile', 'status', 'errors', 'expected_size', 'downloaded',
		     'hint', 'child', '_final_total_size', 'aborted_by_user']

	def __init__(self, url, hint = None):
		"""Create a new download object.
		@param url: the resource to download
		@param hint: object with which this download is associated (an optional hint for the GUI)
		@postcondition: L{status} == L{download_starting}."""
		self.url = url
		self.status = download_starting
		self.hint = hint
		self.aborted_by_user = False

		self.tempfile = None		# Stream for result
		self.errors = None
		self.downloaded = None

		self.expected_size = None	# Final size (excluding skipped bytes)
		self._final_total_size = None	# Set when download is finished

		self.child = None

	def start(self):
		"""Create a temporary file and begin the download.
		@precondition: L{status} == L{download_starting}"""
		assert self.status == download_starting
		assert self.downloaded is None

		self.tempfile = tempfile.TemporaryFile(prefix = 'injector-dl-data-')

		task = tasks.Task(self._do_download(), "download " + self.url)
		self.downloaded = task.finished

	def _do_download(self):
		"""Will trigger L{downloaded} when done (on success or failure)."""
		self.errors = ''

		# Can't use fork here, because Windows doesn't have it
		assert self.child is None, self.child
		self.child = subprocess.Popen([sys.executable, __file__, self.url], stderr = subprocess.PIPE, stdout = self.tempfile)

		self.status = download_fetching

		# Wait for child to exit, collecting error output as we go

		while True:
			yield tasks.InputBlocker(self.child.stderr, "read data from " + self.url)

			data = os.read(self.child.stderr.fileno(), 100)
			if not data:
				break
			self.errors += data

		# Download is complete...

		assert self.status is download_fetching
		assert self.tempfile is not None
		assert self.child is not None

		status = self.child.wait()
		self.child = None

		errors = self.errors
		self.errors = None

		if status and not self.aborted_by_user and not errors:
			errors = 'Download process exited with error status ' \
				 'code ' + hex(status)

		self._final_total_size = self.get_bytes_downloaded_so_far()

		stream = self.tempfile
		self.tempfile = None

		try:
			if self.aborted_by_user:
				raise DownloadAborted(errors)

			if errors:
				raise DownloadError(errors.strip())

			# Check that the download has the correct size, if we know what it should be.
			if self.expected_size is not None:
				size = os.fstat(stream.fileno()).st_size
				if size != self.expected_size:
					raise SafeException('Downloaded archive has incorrect size.\n'
							'URL: %s\n'
							'Expected: %d bytes\n'
							'Received: %d bytes' % (self.url, self.expected_size, size))
		except:
			self.status = download_failed
			_, ex, tb = sys.exc_info()
			self.downloaded.trigger(exception = (ex, tb))
		else:
			self.status = download_complete
			self.downloaded.trigger()

	def abort(self):
		"""Signal the current download to stop.
		@postcondition: L{aborted_by_user}"""
		if self.child is not None:
			info("Killing download process %s", self.child.pid)
			import signal
			os.kill(self.child.pid, signal.SIGTERM)
			self.aborted_by_user = True
		else:
			self.status = download_failed

	def get_current_fraction(self):
		"""Returns the current fraction of this download that has been fetched (from 0 to 1),
		or None if the total size isn't known.
		@return: fraction downloaded
		@rtype: int | None"""
		if self.status is download_starting:
			return 0
		if self.tempfile is None:
			return 1
		if self.expected_size is None:
			return None		# Unknown
		current_size = self.get_bytes_downloaded_so_far()
		return float(current_size) / self.expected_size

	def get_bytes_downloaded_so_far(self):
		"""Get the download progress. Will be zero if the download has not yet started.
		@rtype: int"""
		if self.status is download_starting:
			return 0
		elif self.status is download_fetching:
			return os.fstat(self.tempfile.fileno()).st_size
		else:
			return self._final_total_size

	def __str__(self):
		return "<Download from %s>" % self.url

if __name__ == '__main__':
	def _download_as_child(url):
		from httplib import HTTPException
		from urllib2 import urlopen, HTTPError, URLError
		try:
			#print "Child downloading", url
			if url.startswith('/'):
				if not os.path.isfile(url):
					print >>sys.stderr, "File '%s' does not " \
						"exist!" % url
					return
				src = file(url)
			elif url.startswith('http:') or url.startswith('ftp:'):
				src = urlopen(url)
			else:
				raise Exception('Unsupported URL protocol in: ' + url)

			if hasattr(src, 'fileno'):
				try:
					fd = src.fileno()
				except AttributeError, ex:
					# Hack: fileno on sockets broken in Python 2.4 and 2.5
					# http://bugs.python.org/issue1327971
					fd = src.fp._sock.fp.fileno()

				while True:
					data = os.read(fd, 256)
					if not data: break
					os.write(1, data)
			else:
				# Windows doesn't have fileno
				while True:
					data = src.read(1)
					if not data: break
					os.write(1, data)

			sys.exit(0)
		except (HTTPError, URLError, HTTPException), ex:
			print >>sys.stderr, "Error downloading '" + url + "': " + (str(ex) or str(ex.__class__.__name__))
			sys.exit(1)
	assert len(sys.argv) == 2, "Usage: download URL, not %s" % sys.argv
	_download_as_child(sys.argv[1])
<commit_msg>Don't corrupt downloads on Windows<commit_after>"""
Handles URL downloads.

This is the low-level interface for downloading interfaces, implementations, icons, etc.

@see: L{fetch} higher-level API for downloads that uses this module
"""

# Copyright (C) 2006, Thomas Leonard
# See the README file for details, or visit http://0install.net.

import tempfile, os, sys, subprocess
from zeroinstall import SafeException
from zeroinstall.support import tasks
from logging import info

download_starting = "starting"	# Waiting for UI to start it
download_fetching = "fetching"	# In progress
download_complete = "complete"	# Downloaded and cached OK
download_failed = "failed"

class DownloadError(SafeException):
	"""Download process failed."""
	pass

class DownloadAborted(DownloadError):
	"""Download aborted because of a call to L{Download.abort}"""
	def __init__(self, message):
		SafeException.__init__(self, message or "Download aborted at user's request")

class Download(object):
	"""A download of a single resource to a temporary file.
	@ivar url: the URL of the resource being fetched
	@type url: str
	@ivar tempfile: the file storing the downloaded data
	@type tempfile: file
	@ivar status: the status of the download
	@type status: (download_starting | download_fetching | download_failed | download_complete)
	@ivar errors: data received from the child's stderr
	@type errors: str
	@ivar expected_size: the expected final size of the file
	@type expected_size: int | None
	@ivar downloaded: triggered when the download ends (on success or failure)
	@type downloaded: L{tasks.Blocker}
	@ivar hint: hint passed by and for caller
	@type hint: object
	@ivar child: the child process
	@type child: L{subprocess.Popen}
	@ivar aborted_by_user: whether anyone has called L{abort}
	@type aborted_by_user: bool
	"""
	__slots__ = ['url', 'tempfile', 'status', 'errors', 'expected_size', 'downloaded',
		     'hint', 'child', '_final_total_size', 'aborted_by_user']

	def __init__(self, url, hint = None):
		"""Create a new download object.
		@param url: the resource to download
		@param hint: object with which this download is associated (an optional hint for the GUI)
		@postcondition: L{status} == L{download_starting}."""
		self.url = url
		self.status = download_starting
		self.hint = hint
		self.aborted_by_user = False

		self.tempfile = None		# Stream for result
		self.errors = None
		self.downloaded = None

		self.expected_size = None	# Final size (excluding skipped bytes)
		self._final_total_size = None	# Set when download is finished

		self.child = None

	def start(self):
		"""Create a temporary file and begin the download.
		@precondition: L{status} == L{download_starting}"""
		assert self.status == download_starting
		assert self.downloaded is None

		self.tempfile = tempfile.TemporaryFile(prefix = 'injector-dl-data-')

		task = tasks.Task(self._do_download(), "download " + self.url)
		self.downloaded = task.finished

	def _do_download(self):
		"""Will trigger L{downloaded} when done (on success or failure)."""
		self.errors = ''

		# Can't use fork here, because Windows doesn't have it
		assert self.child is None, self.child
		self.child = subprocess.Popen([sys.executable, '-u', __file__, self.url], stderr = subprocess.PIPE, stdout = self.tempfile)

		self.status = download_fetching

		# Wait for child to exit, collecting error output as we go

		while True:
			yield tasks.InputBlocker(self.child.stderr, "read data from " + self.url)

			data = os.read(self.child.stderr.fileno(), 100)
			if not data:
				break
			self.errors += data

		# Download is complete...

		assert self.status is download_fetching
		assert self.tempfile is not None
		assert self.child is not None

		status = self.child.wait()
		self.child = None

		errors = self.errors
		self.errors = None

		if status and not self.aborted_by_user and not errors:
			errors = 'Download process exited with error status ' \
				 'code ' + hex(status)

		self._final_total_size = self.get_bytes_downloaded_so_far()

		stream = self.tempfile
		self.tempfile = None

		try:
			if self.aborted_by_user:
				raise DownloadAborted(errors)

			if errors:
				raise DownloadError(errors.strip())

			# Check that the download has the correct size, if we know what it should be.
			if self.expected_size is not None:
				size = os.fstat(stream.fileno()).st_size
				if size != self.expected_size:
					raise SafeException('Downloaded archive has incorrect size.\n'
							'URL: %s\n'
							'Expected: %d bytes\n'
							'Received: %d bytes' % (self.url, self.expected_size, size))
		except:
			self.status = download_failed
			_, ex, tb = sys.exc_info()
			self.downloaded.trigger(exception = (ex, tb))
		else:
			self.status = download_complete
			self.downloaded.trigger()

	def abort(self):
		"""Signal the current download to stop.
		@postcondition: L{aborted_by_user}"""
		if self.child is not None:
			info("Killing download process %s", self.child.pid)
			import signal
			os.kill(self.child.pid, signal.SIGTERM)
			self.aborted_by_user = True
		else:
			self.status = download_failed

	def get_current_fraction(self):
		"""Returns the current fraction of this download that has been fetched (from 0 to 1),
		or None if the total size isn't known.
		@return: fraction downloaded
		@rtype: int | None"""
		if self.status is download_starting:
			return 0
		if self.tempfile is None:
			return 1
		if self.expected_size is None:
			return None		# Unknown
		current_size = self.get_bytes_downloaded_so_far()
		return float(current_size) / self.expected_size

	def get_bytes_downloaded_so_far(self):
		"""Get the download progress. Will be zero if the download has not yet started.
		@rtype: int"""
		if self.status is download_starting:
			return 0
		elif self.status is download_fetching:
			return os.fstat(self.tempfile.fileno()).st_size
		else:
			return self._final_total_size

	def __str__(self):
		return "<Download from %s>" % self.url

if __name__ == '__main__':
	def _download_as_child(url):
		from httplib import HTTPException
		from urllib2 import urlopen, HTTPError, URLError
		try:
			#print "Child downloading", url
			if url.startswith('/'):
				if not os.path.isfile(url):
					print >>sys.stderr, "File '%s' does not " \
						"exist!" % url
					return
				src = file(url)
			elif url.startswith('http:') or url.startswith('ftp:'):
				src = urlopen(url)
			else:
				raise Exception('Unsupported URL protocol in: ' + url)

			if hasattr(src, 'fileno'):
				try:
					fd = src.fileno()
				except AttributeError, ex:
					# Hack: fileno on sockets broken in Python 2.4 and 2.5
					# http://bugs.python.org/issue1327971
					fd = src.fp._sock.fp.fileno()

				while True:
					data = os.read(fd, 256)
					if not data: break
					os.write(1, data)
			else:
				# Windows doesn't have fileno
				while True:
					data = src.read(1)
					if not data: break
					os.write(1, data)

			sys.exit(0)
		except (HTTPError, URLError, HTTPException), ex:
			print >>sys.stderr, "Error downloading '" + url + "': " + (str(ex) or str(ex.__class__.__name__))
			sys.exit(1)
	assert len(sys.argv) == 2, "Usage: download URL, not %s" % sys.argv
	_download_as_child(sys.argv[1])
<|endoftext|>


<commit_before>from __future__ import absolute_import

import random
import requests
from typing import Any, Dict, List, Optional, SupportsInt, Text, Union, Type

from version import ZULIP_VERSION
from zerver.models import PushDeviceToken, Message, Recipient, UserProfile, \
    UserMessage, get_display_recipient, receives_offline_notifications, \
    receives_online_notifications
from zerver.models import get_user_profile_by_id
from zerver.lib.avatar import absolute_avatar_url
from zerver.lib.request import JsonableError
from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
from zerver.decorator import statsd_increment
from zerver.lib.utils import generate_random_token
from zerver.lib.queue import retry_event

from gcm import GCM

from django.conf import settings
from django.utils.timezone import now as timezone_now
from django.utils.translation import ugettext as _
from six.moves import urllib

import base64
import binascii
import logging
import os
import time
import ujson
from functools import partial

if settings.ZILENCER_ENABLED:
    from zilencer.models import RemotePushDeviceToken
else:  # nocoverage  -- Not convenient to add test for this.
    from mock import Mock
    RemotePushDeviceToken = Mock()  # type: ignore # https://github.com/JukkaL/mypy/issues/1188

DeviceToken = Union[PushDeviceToken, RemotePushDeviceToken]

# We store the token as b64, but apns-client wants hex strings
def b64_to_hex(data):
    # type: (bytes) -> Text
    return binascii.hexlify(base64.b64decode(data)).decode('utf-8')

def hex_to_b64(data):
    # type: (Text) -> bytes
    return base64.b64encode(binascii.unhexlify(data.encode('utf-8')))

#
# Sending to APNs, for iOS
#

# `APNS_SANDBOX` should be a bool
assert isinstance(settings.APNS_SANDBOX, bool)

@statsd_increment("apple_push_notification")
def send_apple_push_notification(user_id, devices, **extra_data):
    # type: (int, List[DeviceToken], **Any) -> None
    if not devices:
        return
    logging.warn("APNs unimplemented.  Dropping notification for user %d with %d devices.",
                 user_id, len(devices))

#
# Sending to GCM, for Android
#

if settings.ANDROID_GCM_API_KEY:  # nocoverage
    gcm = GCM(settings.ANDROID_GCM_API_KEY)
else:
    gcm = None

def send_android_push_notification_to_user(user_profile, data):
    # type: (UserProfile, Dict[str, Any]) -> None
    devices = list(PushDeviceToken.objects.filter(user=user_profile,
                                                  kind=PushDeviceToken.GCM))
    send_android_push_notification(devices, data)

@statsd_increment("android_push_notification")
def send_android_push_notification(devices, data, remote=False):
    # type: (List[DeviceToken], Dict[str, Any], bool) -> None
    if not gcm:
        logging.warning("Attempting to send a GCM push notification, but no API key was configured")
        return
    reg_ids = [device.token for device in devices]

    if remote:
        DeviceTokenClass = RemotePushDeviceToken
    else:
        DeviceTokenClass = PushDeviceToken

    try:
        res = gcm.json_request(registration_ids=reg_ids, data=data, retries=10)
    except IOError as e:
        logging.warning(str(e))
        return

    if res and 'success' in res:
        for reg_id, msg_id in res['success'].items():
            logging.info("GCM: Sent %s as %s" % (reg_id, msg_id))

    # res.canonical will contain results when there are duplicate registrations for the same
    # device. The "canonical" registration is the latest registration made by the device.
    # Ref: http://developer.android.com/google/gcm/adv.html#canonical
    if 'canonical' in res:
        for reg_id, new_reg_id in res['canonical'].items():
            if reg_id == new_reg_id:
                # I'm not sure if this should happen. In any case, not really actionable.
                logging.warning("GCM: Got canonical ref but it already matches our ID %s!" % (reg_id,))
            elif not DeviceTokenClass.objects.filter(token=new_reg_id,
                                                     kind=DeviceTokenClass.GCM).count():
                # This case shouldn't happen; any time we get a canonical ref it should have been
                # previously registered in our system.
                #
                # That said, recovery is easy: just update the current PDT object to use the new ID.
                logging.warning(
                    "GCM: Got canonical ref %s replacing %s but new ID not registered! Updating." %
                    (new_reg_id, reg_id))
                DeviceTokenClass.objects.filter(
                    token=reg_id, kind=DeviceTokenClass.GCM).update(token=new_reg_id)
            else:
                # Since we know the new ID is registered in our system we can just drop the old one.
                logging.info("GCM: Got canonical ref %s, dropping %s" % (new_reg_id, reg_id))

                DeviceTokenClass.objects.filter(token=reg_id, kind=DeviceTokenClass.GCM).delete()

    if 'errors' in res:
        for error, reg_ids in res['errors'].items():
            if error in ['NotRegistered', 'InvalidRegistration']:
                for reg_id in reg_ids:
                    logging.info("GCM: Removing %s" % (reg_id,))

                    device = DeviceTokenClass.objects.get(token=reg_id, kind=DeviceTokenClass.GCM)
                    device.delete()
            else:
                for reg_id in reg_ids:
                    logging.warning("GCM: Delivery to %s failed: %s" % (reg_id, error))

    # python-gcm handles retrying of the unsent messages.
    # Ref: https://github.com/geeknam/python-gcm/blob/master/gcm/gcm.py#L497

#
# Sending to a bouncer
#

def uses_notification_bouncer():
    # type: () -> bool
    return settings.PUSH_NOTIFICATION_BOUNCER_URL is not None

def send_notifications_to_bouncer(user_profile_id, apns_payload, gcm_payload):
    # type: (int, Dict[str, Any], Dict[str, Any]) -> None
    post_data = {
        'user_id': user_profile_id,
        'apns_payload': apns_payload,
        'gcm_payload': gcm_payload,
    }
    send_json_to_push_bouncer('POST', 'notify', post_data)

def send_json_to_push_bouncer(method, endpoint, post_data):
    # type: (str, str, Dict[str, Any]) -> None
    send_to_push_bouncer(
        method,
        endpoint,
        ujson.dumps(post_data),
        extra_headers={"Content-type": "application/json"},
    )

def send_to_push_bouncer(method, endpoint, post_data, extra_headers=None):
    # type: (str, str, Union[Text, Dict[str, Any]], Optional[Dict[str, Any]]) -> None
    url = urllib.parse.urljoin(settings.PUSH_NOTIFICATION_BOUNCER_URL,
                               '/api/v1/remotes/push/' + endpoint)
    api_auth = requests.auth.HTTPBasicAuth(settings.ZULIP_ORG_ID,
                                           settings.ZULIP_ORG_KEY)

    headers = {"User-agent": "ZulipServer/%s" % (ZULIP_VERSION,)}
    if extra_headers is not None:
        headers.update(extra_headers)

    res = requests.request(method,
                           url,
                           data=post_data,
                           auth=api_auth,
                           timeout=30,
                           verify=True,
                           headers=headers)

    # TODO: Think more carefully about how this error hanlding should work.
    if res.status_code >= 500:
        raise JsonableError(_("Error received from push notification bouncer"))
    elif res.status_code >= 400:
        try:
            msg = ujson.loads(res.content)['msg']
        except Exception:
            raise JsonableError(_("Error received from push notification bouncer"))
        raise JsonableError(msg)
    elif res.status_code != 200:
        raise JsonableError(_("Error received from push notification bouncer"))

    # If we don't throw an exception, it's a successful bounce!

#
# Managing device tokens
#

def num_push_devices_for_user(user_profile, kind = None):
    # type: (UserProfile, Optional[int]) -> PushDeviceToken
    if kind is None:
        return PushDeviceToken.objects.filter(user=user_profile).count()
    else:
        return PushDeviceToken.objects.filter(user=user_profile, kind=kind).count()

def add_push_device_token(user_profile, token_str, kind, ios_app_id=None):
    # type: (UserProfile, bytes, int, Optional[str]) -> None

    logging.info("New push device: %d %r %d %r",
                 user_profile.id, token_str, kind, ios_app_id)

    # If we're sending things to the push notification bouncer
    # register this user with them here
    if uses_notification_bouncer():
        post_data = {
            'server_uuid': settings.ZULIP_ORG_ID,
            'user_id': user_profile.id,
            'token': token_str,
            'token_kind': kind,
        }

        if kind == PushDeviceToken.APNS:
            post_data['ios_app_id'] = ios_app_id

        logging.info("Sending new push device to bouncer: %r", post_data)
        send_to_push_bouncer('POST', 'register', post_data)
        return

    # If another user was previously logged in on the same device and didn't
    # properly log out, the token will still be registered to the wrong account
    PushDeviceToken.objects.filter(token=token_str).exclude(user=user_profile).delete()

    # Overwrite with the latest value
    token, created = PushDeviceToken.objects.get_or_create(user=user_profile,
                                                           token=token_str,
                                                           defaults=dict(
                                                               kind=kind,
                                                               ios_app_id=ios_app_id))
    if not created:
        logging.info("Existing push device updated.")
        token.last_updated = timezone_now()
        token.save(update_fields=['last_updated'])
    else:
        logging.info("New push device created.")

def remove_push_device_token(user_profile, token_str, kind):
    # type: (UserProfile, bytes, int) -> None

    # If we're sending things to the push notification bouncer
    # register this user with them here
    if uses_notification_bouncer():
        # TODO: Make this a remove item
        post_data = {
            'server_uuid': settings.ZULIP_ORG_ID,
            'user_id': user_profile.id,
            'token': token_str,
            'token_kind': kind,
        }
        send_to_push_bouncer("POST", "unregister", post_data)
        return

    try:
        token = PushDeviceToken.objects.get(token=token_str, kind=kind)
        token.delete()
    except PushDeviceToken.DoesNotExist:
        raise JsonableError(_("Token does not exist"))

#
# Push notifications in general
#

def get_alert_from_message(message):
    # type: (Message) -> Text
    """
    Determine what alert string to display based on the missed messages.
    """
    sender_str = message.sender.full_name
    if message.recipient.type == Recipient.HUDDLE:
        return "New private group message from %s" % (sender_str,)
    elif message.recipient.type == Recipient.PERSONAL:
        return "New private message from %s" % (sender_str,)
    elif message.recipient.type == Recipient.STREAM:
        return "New mention from %s" % (sender_str,)
    else:
        return "New Zulip mentions and private messages from %s" % (sender_str,)

def get_apns_payload(message):
    # type: (Message) -> Dict[str, Any]
    return {
        'alert': get_alert_from_message(message),
        'message_ids': [message.id],
    }

def get_gcm_payload(user_profile, message):
    # type: (UserProfile, Message) -> Dict[str, Any]
    content = message.content
    content_truncated = (len(content) > 200)
    if content_truncated:
        content = content[:200] + "..."

    android_data = {
        'user': user_profile.email,
        'event': 'message',
        'alert': get_alert_from_message(message),
        'zulip_message_id': message.id,  # message_id is reserved for CCS
        'time': datetime_to_timestamp(message.pub_date),
        'content': content,
        'content_truncated': content_truncated,
        'sender_email': message.sender.email,
        'sender_full_name': message.sender.full_name,
        'sender_avatar_url': absolute_avatar_url(message.sender),
    }

    if message.recipient.type == Recipient.STREAM:
        android_data['recipient_type'] = "stream"
        android_data['stream'] = get_display_recipient(message.recipient)
        android_data['topic'] = message.subject
    elif message.recipient.type in (Recipient.HUDDLE, Recipient.PERSONAL):
        android_data['recipient_type'] = "private"

    return android_data

@statsd_increment("push_notifications")
def handle_push_notification(user_profile_id, missed_message):
    # type: (int, Dict[str, Any]) -> None
    """
    missed_message is the event received by the
    zerver.worker.queue_processors.PushNotificationWorker.consume function.
    """
    try:
        user_profile = get_user_profile_by_id(user_profile_id)
        if not (receives_offline_notifications(user_profile) or receives_online_notifications(user_profile)):
            return

        umessage = UserMessage.objects.get(user_profile=user_profile,
                                           message__id=missed_message['message_id'])
        message = umessage.message
        if umessage.flags.read:
            return

        apns_payload = get_apns_payload(message)
        gcm_payload = get_gcm_payload(user_profile, message)

        if uses_notification_bouncer():
            try:
                send_notifications_to_bouncer(user_profile_id,
                                              apns_payload,
                                              gcm_payload)
            except requests.ConnectionError:
                if 'failed_tries' not in missed_message:
                    missed_message['failed_tries'] = 0

                def failure_processor(event):
                    # type: (Dict[str, Any]) -> None
                    logging.warning("Maximum retries exceeded for trigger:%s event:push_notification" % (event['user_profile_id']))
                retry_event('missedmessage_mobile_notifications', missed_message,
                            failure_processor)

            return

        android_devices = list(PushDeviceToken.objects.filter(user=user_profile,
                                                              kind=PushDeviceToken.GCM))

        apple_devices = list(PushDeviceToken.objects.filter(user=user_profile,
                                                            kind=PushDeviceToken.APNS))

        # TODO: set badge count in a better way
        if apple_devices:
            send_apple_push_notification(user_profile.id, apple_devices,
                                         badge=1, zulip=apns_payload)

        if android_devices:
            send_android_push_notification(android_devices, gcm_payload)

    except UserMessage.DoesNotExist:
        logging.error("Could not find UserMessage with message_id %s" % (missed_message['message_id'],))
<commit_msg>push notifs: Organize imports.<commit_after>from __future__ import absolute_import

import base64
import binascii
from functools import partial
import logging
import os
import time
import random
from typing import Any, Dict, List, Optional, SupportsInt, Text, Union, Type

from django.conf import settings
from django.utils.timezone import now as timezone_now
from django.utils.translation import ugettext as _
from gcm import GCM
import requests
from six.moves import urllib
import ujson

from zerver.decorator import statsd_increment
from zerver.lib.avatar import absolute_avatar_url
from zerver.lib.queue import retry_event
from zerver.lib.request import JsonableError
from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
from zerver.lib.utils import generate_random_token
from zerver.models import PushDeviceToken, Message, Recipient, UserProfile, \
    UserMessage, get_display_recipient, receives_offline_notifications, \
    receives_online_notifications, get_user_profile_by_id
from version import ZULIP_VERSION

if settings.ZILENCER_ENABLED:
    from zilencer.models import RemotePushDeviceToken
else:  # nocoverage  -- Not convenient to add test for this.
    from mock import Mock
    RemotePushDeviceToken = Mock()  # type: ignore # https://github.com/JukkaL/mypy/issues/1188

DeviceToken = Union[PushDeviceToken, RemotePushDeviceToken]

# We store the token as b64, but apns-client wants hex strings
def b64_to_hex(data):
    # type: (bytes) -> Text
    return binascii.hexlify(base64.b64decode(data)).decode('utf-8')

def hex_to_b64(data):
    # type: (Text) -> bytes
    return base64.b64encode(binascii.unhexlify(data.encode('utf-8')))

#
# Sending to APNs, for iOS
#

# `APNS_SANDBOX` should be a bool
assert isinstance(settings.APNS_SANDBOX, bool)

@statsd_increment("apple_push_notification")
def send_apple_push_notification(user_id, devices, **extra_data):
    # type: (int, List[DeviceToken], **Any) -> None
    if not devices:
        return
    logging.warn("APNs unimplemented.  Dropping notification for user %d with %d devices.",
                 user_id, len(devices))

#
# Sending to GCM, for Android
#

if settings.ANDROID_GCM_API_KEY:  # nocoverage
    gcm = GCM(settings.ANDROID_GCM_API_KEY)
else:
    gcm = None

def send_android_push_notification_to_user(user_profile, data):
    # type: (UserProfile, Dict[str, Any]) -> None
    devices = list(PushDeviceToken.objects.filter(user=user_profile,
                                                  kind=PushDeviceToken.GCM))
    send_android_push_notification(devices, data)

@statsd_increment("android_push_notification")
def send_android_push_notification(devices, data, remote=False):
    # type: (List[DeviceToken], Dict[str, Any], bool) -> None
    if not gcm:
        logging.warning("Attempting to send a GCM push notification, but no API key was configured")
        return
    reg_ids = [device.token for device in devices]

    if remote:
        DeviceTokenClass = RemotePushDeviceToken
    else:
        DeviceTokenClass = PushDeviceToken

    try:
        res = gcm.json_request(registration_ids=reg_ids, data=data, retries=10)
    except IOError as e:
        logging.warning(str(e))
        return

    if res and 'success' in res:
        for reg_id, msg_id in res['success'].items():
            logging.info("GCM: Sent %s as %s" % (reg_id, msg_id))

    # res.canonical will contain results when there are duplicate registrations for the same
    # device. The "canonical" registration is the latest registration made by the device.
    # Ref: http://developer.android.com/google/gcm/adv.html#canonical
    if 'canonical' in res:
        for reg_id, new_reg_id in res['canonical'].items():
            if reg_id == new_reg_id:
                # I'm not sure if this should happen. In any case, not really actionable.
                logging.warning("GCM: Got canonical ref but it already matches our ID %s!" % (reg_id,))
            elif not DeviceTokenClass.objects.filter(token=new_reg_id,
                                                     kind=DeviceTokenClass.GCM).count():
                # This case shouldn't happen; any time we get a canonical ref it should have been
                # previously registered in our system.
                #
                # That said, recovery is easy: just update the current PDT object to use the new ID.
                logging.warning(
                    "GCM: Got canonical ref %s replacing %s but new ID not registered! Updating." %
                    (new_reg_id, reg_id))
                DeviceTokenClass.objects.filter(
                    token=reg_id, kind=DeviceTokenClass.GCM).update(token=new_reg_id)
            else:
                # Since we know the new ID is registered in our system we can just drop the old one.
                logging.info("GCM: Got canonical ref %s, dropping %s" % (new_reg_id, reg_id))

                DeviceTokenClass.objects.filter(token=reg_id, kind=DeviceTokenClass.GCM).delete()

    if 'errors' in res:
        for error, reg_ids in res['errors'].items():
            if error in ['NotRegistered', 'InvalidRegistration']:
                for reg_id in reg_ids:
                    logging.info("GCM: Removing %s" % (reg_id,))

                    device = DeviceTokenClass.objects.get(token=reg_id, kind=DeviceTokenClass.GCM)
                    device.delete()
            else:
                for reg_id in reg_ids:
                    logging.warning("GCM: Delivery to %s failed: %s" % (reg_id, error))

    # python-gcm handles retrying of the unsent messages.
    # Ref: https://github.com/geeknam/python-gcm/blob/master/gcm/gcm.py#L497

#
# Sending to a bouncer
#

def uses_notification_bouncer():
    # type: () -> bool
    return settings.PUSH_NOTIFICATION_BOUNCER_URL is not None

def send_notifications_to_bouncer(user_profile_id, apns_payload, gcm_payload):
    # type: (int, Dict[str, Any], Dict[str, Any]) -> None
    post_data = {
        'user_id': user_profile_id,
        'apns_payload': apns_payload,
        'gcm_payload': gcm_payload,
    }
    send_json_to_push_bouncer('POST', 'notify', post_data)

def send_json_to_push_bouncer(method, endpoint, post_data):
    # type: (str, str, Dict[str, Any]) -> None
    send_to_push_bouncer(
        method,
        endpoint,
        ujson.dumps(post_data),
        extra_headers={"Content-type": "application/json"},
    )

def send_to_push_bouncer(method, endpoint, post_data, extra_headers=None):
    # type: (str, str, Union[Text, Dict[str, Any]], Optional[Dict[str, Any]]) -> None
    url = urllib.parse.urljoin(settings.PUSH_NOTIFICATION_BOUNCER_URL,
                               '/api/v1/remotes/push/' + endpoint)
    api_auth = requests.auth.HTTPBasicAuth(settings.ZULIP_ORG_ID,
                                           settings.ZULIP_ORG_KEY)

    headers = {"User-agent": "ZulipServer/%s" % (ZULIP_VERSION,)}
    if extra_headers is not None:
        headers.update(extra_headers)

    res = requests.request(method,
                           url,
                           data=post_data,
                           auth=api_auth,
                           timeout=30,
                           verify=True,
                           headers=headers)

    # TODO: Think more carefully about how this error hanlding should work.
    if res.status_code >= 500:
        raise JsonableError(_("Error received from push notification bouncer"))
    elif res.status_code >= 400:
        try:
            msg = ujson.loads(res.content)['msg']
        except Exception:
            raise JsonableError(_("Error received from push notification bouncer"))
        raise JsonableError(msg)
    elif res.status_code != 200:
        raise JsonableError(_("Error received from push notification bouncer"))

    # If we don't throw an exception, it's a successful bounce!

#
# Managing device tokens
#

def num_push_devices_for_user(user_profile, kind = None):
    # type: (UserProfile, Optional[int]) -> PushDeviceToken
    if kind is None:
        return PushDeviceToken.objects.filter(user=user_profile).count()
    else:
        return PushDeviceToken.objects.filter(user=user_profile, kind=kind).count()

def add_push_device_token(user_profile, token_str, kind, ios_app_id=None):
    # type: (UserProfile, bytes, int, Optional[str]) -> None

    logging.info("New push device: %d %r %d %r",
                 user_profile.id, token_str, kind, ios_app_id)

    # If we're sending things to the push notification bouncer
    # register this user with them here
    if uses_notification_bouncer():
        post_data = {
            'server_uuid': settings.ZULIP_ORG_ID,
            'user_id': user_profile.id,
            'token': token_str,
            'token_kind': kind,
        }

        if kind == PushDeviceToken.APNS:
            post_data['ios_app_id'] = ios_app_id

        logging.info("Sending new push device to bouncer: %r", post_data)
        send_to_push_bouncer('POST', 'register', post_data)
        return

    # If another user was previously logged in on the same device and didn't
    # properly log out, the token will still be registered to the wrong account
    PushDeviceToken.objects.filter(token=token_str).exclude(user=user_profile).delete()

    # Overwrite with the latest value
    token, created = PushDeviceToken.objects.get_or_create(user=user_profile,
                                                           token=token_str,
                                                           defaults=dict(
                                                               kind=kind,
                                                               ios_app_id=ios_app_id))
    if not created:
        logging.info("Existing push device updated.")
        token.last_updated = timezone_now()
        token.save(update_fields=['last_updated'])
    else:
        logging.info("New push device created.")

def remove_push_device_token(user_profile, token_str, kind):
    # type: (UserProfile, bytes, int) -> None

    # If we're sending things to the push notification bouncer
    # register this user with them here
    if uses_notification_bouncer():
        # TODO: Make this a remove item
        post_data = {
            'server_uuid': settings.ZULIP_ORG_ID,
            'user_id': user_profile.id,
            'token': token_str,
            'token_kind': kind,
        }
        send_to_push_bouncer("POST", "unregister", post_data)
        return

    try:
        token = PushDeviceToken.objects.get(token=token_str, kind=kind)
        token.delete()
    except PushDeviceToken.DoesNotExist:
        raise JsonableError(_("Token does not exist"))

#
# Push notifications in general
#

def get_alert_from_message(message):
    # type: (Message) -> Text
    """
    Determine what alert string to display based on the missed messages.
    """
    sender_str = message.sender.full_name
    if message.recipient.type == Recipient.HUDDLE:
        return "New private group message from %s" % (sender_str,)
    elif message.recipient.type == Recipient.PERSONAL:
        return "New private message from %s" % (sender_str,)
    elif message.recipient.type == Recipient.STREAM:
        return "New mention from %s" % (sender_str,)
    else:
        return "New Zulip mentions and private messages from %s" % (sender_str,)

def get_apns_payload(message):
    # type: (Message) -> Dict[str, Any]
    return {
        'alert': get_alert_from_message(message),
        'message_ids': [message.id],
    }

def get_gcm_payload(user_profile, message):
    # type: (UserProfile, Message) -> Dict[str, Any]
    content = message.content
    content_truncated = (len(content) > 200)
    if content_truncated:
        content = content[:200] + "..."

    android_data = {
        'user': user_profile.email,
        'event': 'message',
        'alert': get_alert_from_message(message),
        'zulip_message_id': message.id,  # message_id is reserved for CCS
        'time': datetime_to_timestamp(message.pub_date),
        'content': content,
        'content_truncated': content_truncated,
        'sender_email': message.sender.email,
        'sender_full_name': message.sender.full_name,
        'sender_avatar_url': absolute_avatar_url(message.sender),
    }

    if message.recipient.type == Recipient.STREAM:
        android_data['recipient_type'] = "stream"
        android_data['stream'] = get_display_recipient(message.recipient)
        android_data['topic'] = message.subject
    elif message.recipient.type in (Recipient.HUDDLE, Recipient.PERSONAL):
        android_data['recipient_type'] = "private"

    return android_data

@statsd_increment("push_notifications")
def handle_push_notification(user_profile_id, missed_message):
    # type: (int, Dict[str, Any]) -> None
    """
    missed_message is the event received by the
    zerver.worker.queue_processors.PushNotificationWorker.consume function.
    """
    try:
        user_profile = get_user_profile_by_id(user_profile_id)
        if not (receives_offline_notifications(user_profile) or receives_online_notifications(user_profile)):
            return

        umessage = UserMessage.objects.get(user_profile=user_profile,
                                           message__id=missed_message['message_id'])
        message = umessage.message
        if umessage.flags.read:
            return

        apns_payload = get_apns_payload(message)
        gcm_payload = get_gcm_payload(user_profile, message)

        if uses_notification_bouncer():
            try:
                send_notifications_to_bouncer(user_profile_id,
                                              apns_payload,
                                              gcm_payload)
            except requests.ConnectionError:
                if 'failed_tries' not in missed_message:
                    missed_message['failed_tries'] = 0

                def failure_processor(event):
                    # type: (Dict[str, Any]) -> None
                    logging.warning("Maximum retries exceeded for trigger:%s event:push_notification" % (event['user_profile_id']))
                retry_event('missedmessage_mobile_notifications', missed_message,
                            failure_processor)

            return

        android_devices = list(PushDeviceToken.objects.filter(user=user_profile,
                                                              kind=PushDeviceToken.GCM))

        apple_devices = list(PushDeviceToken.objects.filter(user=user_profile,
                                                            kind=PushDeviceToken.APNS))

        # TODO: set badge count in a better way
        if apple_devices:
            send_apple_push_notification(user_profile.id, apple_devices,
                                         badge=1, zulip=apns_payload)

        if android_devices:
            send_android_push_notification(android_devices, gcm_payload)

    except UserMessage.DoesNotExist:
        logging.error("Could not find UserMessage with message_id %s" % (missed_message['message_id'],))
<|endoftext|>


<commit_before>from zerver.lib.alert_words import (
    add_user_alert_words,
    alert_words_in_realm,
    remove_user_alert_words,
    user_alert_words,
)

from zerver.lib.test_helpers import (
    most_recent_message,
    most_recent_usermessage,
)

from zerver.lib.test_classes import (
    ZulipTestCase,
)

from zerver.models import (
    UserProfile,
)

import ujson

class AlertWordTests(ZulipTestCase):
    interesting_alert_word_list = ['alert', 'multi-word word', '☃']

    def test_internal_endpoint(self) -> None:
        user_name = "cordelia"
        self.login(user_name)

        params = {
            'alert_words': ujson.dumps(['milk', 'cookies'])
        }
        result = self.client_post('/json/users/me/alert_words', params)
        self.assert_json_success(result)
        user = self.example_user(user_name)
        words = user_alert_words(user)
        self.assertEqual(set(words), {'milk', 'cookies'})

    def test_default_no_words(self) -> None:
        """
        Users start out with no alert words.
        """
        user = self.example_user('cordelia')
        words = user_alert_words(user)
        self.assertEqual(words, [])

    def test_basics(self) -> None:
        """
        Verifies the basic behavior of modifying alert words.

        Also verifies the cache-flushing behavior.
        """
        user = self.example_user('cordelia')
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words.get(user.id, []), 0)

        # Add several words, including multi-word and non-ascii words.
        add_user_alert_words(user, self.interesting_alert_word_list)

        words = user_alert_words(user)
        self.assertEqual(set(words), set(self.interesting_alert_word_list))
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words[user.id], 3)

        # Test the case-insensitivity of adding words
        add_user_alert_words(user, set(["ALert", "ALERT"]))
        words = user_alert_words(user)
        self.assertEqual(set(words), set(self.interesting_alert_word_list))
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words[user.id], 3)

        # Test the case-insensitivity of removing words
        remove_user_alert_words(user, set(["ALert"]))
        words = user_alert_words(user)
        self.assertEqual(set(words), set(self.interesting_alert_word_list) - {'alert'})
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words[user.id], 2)

    def test_remove_word(self) -> None:
        """
        Removing alert words works via remove_user_alert_words, even
        for multi-word and non-ascii words.
        """
        user = self.example_user('cordelia')

        expected_remaining_alerts = set(self.interesting_alert_word_list)
        add_user_alert_words(user, self.interesting_alert_word_list)

        for alert_word in self.interesting_alert_word_list:
            remove_user_alert_words(user, [alert_word])
            expected_remaining_alerts.remove(alert_word)
            actual_remaining_alerts = user_alert_words(user)
            self.assertEqual(set(actual_remaining_alerts),
                             expected_remaining_alerts)

    def test_realm_words(self) -> None:
        """
        We can gather alert words for an entire realm via
        alert_words_in_realm. Alerts added for one user do not impact other
        users.
        """
        user1 = self.example_user('cordelia')

        add_user_alert_words(user1, self.interesting_alert_word_list)

        user2 = self.example_user('othello')
        add_user_alert_words(user2, ['another'])

        realm_words = alert_words_in_realm(user2.realm)
        self.assertEqual(len(realm_words), 2)
        self.assertEqual(set(realm_words.keys()), {user1.id, user2.id})
        self.assertEqual(set(realm_words[user1.id]),
                         set(self.interesting_alert_word_list))
        self.assertEqual(set(realm_words[user2.id]), {'another'})

    def test_json_list_default(self) -> None:
        self.login('hamlet')

        result = self.client_get('/json/users/me/alert_words')
        self.assert_json_success(result)
        self.assertEqual(result.json()['alert_words'], [])

    def test_json_list_nonempty(self) -> None:
        hamlet = self.example_user('hamlet')
        add_user_alert_words(hamlet, ['one', 'two', 'three'])

        self.login('hamlet')
        result = self.client_get('/json/users/me/alert_words')
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

    def test_json_list_add(self) -> None:
        self.login('hamlet')

        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['one ', '\n two', 'three'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

    def test_json_list_remove(self) -> None:
        self.login('hamlet')

        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['one', 'two', 'three'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

        result = self.client_delete('/json/users/me/alert_words',
                                    {'alert_words': ujson.dumps(['one'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'two', 'three'})

    def message_does_alert(self, user_profile: UserProfile, message: str) -> bool:
        """Send a bunch of messages as othello, so Hamlet is notified"""
        self.send_stream_message(self.example_user("othello"), "Denmark", message)
        user_message = most_recent_usermessage(user_profile)
        return 'has_alert_word' in user_message.flags_list()

    def test_alert_flags(self) -> None:
        self.login('hamlet')
        user_profile_hamlet = self.example_user('hamlet')

        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['one', 'two', 'three'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

        # Alerts in the middle of messages work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Normal alert one time"))
        # Alerts at the end of messages work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Normal alert one"))
        # Alerts at the beginning of messages work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "two normal alerts"))
        # Alerts with surrounding punctuation work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "This one? should alert"))
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Definitely time for three."))
        # Multiple alerts in a message work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "One two three o'clock"))
        # Alerts are case-insensitive.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "One o'clock"))
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Case of ONE, won't stop me"))

        # We don't cause alerts for matches in URLs.
        self.assertFalse(self.message_does_alert(user_profile_hamlet, "Don't alert on http://t.co/one/ urls"))
        self.assertFalse(self.message_does_alert(user_profile_hamlet, "Don't alert on http://t.co/one urls"))

    def test_update_alert_words(self) -> None:
        user_profile = self.example_user('hamlet')

        self.login_user(user_profile)
        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['ALERT'])})

        content = 'this is an ALERT for you'
        self.send_stream_message(user_profile, "Denmark", content)
        self.assert_json_success(result)

        original_message = most_recent_message(user_profile)

        user_message = most_recent_usermessage(user_profile)
        self.assertIn('has_alert_word', user_message.flags_list())

        result = self.client_patch("/json/messages/" + str(original_message.id), {
            'message_id': original_message.id,
            'content': 'new ALERT for you',
        })
        self.assert_json_success(result)

        user_message = most_recent_usermessage(user_profile)
        self.assertEqual(user_message.message.content, 'new ALERT for you')
        self.assertIn('has_alert_word', user_message.flags_list())

        result = self.client_patch("/json/messages/" + str(original_message.id), {
            'message_id': original_message.id,
            'content': 'sorry false alarm',
        })
        self.assert_json_success(result)

        user_message = most_recent_usermessage(user_profile)
        self.assertEqual(user_message.message.content, 'sorry false alarm')
        self.assertNotIn('has_alert_word', user_message.flags_list())
<commit_msg>alert_words: Change test_alert_words to use actions.do_*().<commit_after>from zerver.lib.alert_words import (
    alert_words_in_realm,
    user_alert_words,
)

from zerver.lib.actions import (
    do_add_alert_words,
    do_remove_alert_words,
)

from zerver.lib.test_helpers import (
    most_recent_message,
    most_recent_usermessage,
)

from zerver.lib.test_classes import (
    ZulipTestCase,
)

from zerver.models import (
    UserProfile,
)

import ujson

class AlertWordTests(ZulipTestCase):
    interesting_alert_word_list = ['alert', 'multi-word word', '☃']

    def test_internal_endpoint(self) -> None:
        user_name = "cordelia"
        self.login(user_name)

        params = {
            'alert_words': ujson.dumps(['milk', 'cookies'])
        }
        result = self.client_post('/json/users/me/alert_words', params)
        self.assert_json_success(result)
        user = self.example_user(user_name)
        words = user_alert_words(user)
        self.assertEqual(set(words), {'milk', 'cookies'})

    def test_default_no_words(self) -> None:
        """
        Users start out with no alert words.
        """
        user = self.example_user('cordelia')
        words = user_alert_words(user)
        self.assertEqual(words, [])

    def test_basics(self) -> None:
        """
        Verifies the basic behavior of modifying alert words.

        Also verifies the cache-flushing behavior.
        """
        user = self.example_user('cordelia')
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words.get(user.id, []), 0)

        # Add several words, including multi-word and non-ascii words.
        do_add_alert_words(user, self.interesting_alert_word_list)

        words = user_alert_words(user)
        self.assertEqual(set(words), set(self.interesting_alert_word_list))
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words[user.id], 3)

        # Test the case-insensitivity of adding words
        do_add_alert_words(user, set(["ALert", "ALERT"]))
        words = user_alert_words(user)
        self.assertEqual(set(words), set(self.interesting_alert_word_list))
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words[user.id], 3)

        # Test the case-insensitivity of removing words
        do_remove_alert_words(user, set(["ALert"]))
        words = user_alert_words(user)
        self.assertEqual(set(words), set(self.interesting_alert_word_list) - {'alert'})
        realm_alert_words = alert_words_in_realm(user.realm)
        self.assert_length(realm_alert_words[user.id], 2)

    def test_remove_word(self) -> None:
        """
        Removing alert words works via do_remove_alert_words, even
        for multi-word and non-ascii words.
        """
        user = self.example_user('cordelia')

        expected_remaining_alerts = set(self.interesting_alert_word_list)
        do_add_alert_words(user, self.interesting_alert_word_list)

        for alert_word in self.interesting_alert_word_list:
            do_remove_alert_words(user, [alert_word])
            expected_remaining_alerts.remove(alert_word)
            actual_remaining_alerts = user_alert_words(user)
            self.assertEqual(set(actual_remaining_alerts),
                             expected_remaining_alerts)

    def test_realm_words(self) -> None:
        """
        We can gather alert words for an entire realm via
        alert_words_in_realm. Alerts added for one user do not impact other
        users.
        """
        user1 = self.example_user('cordelia')

        do_add_alert_words(user1, self.interesting_alert_word_list)

        user2 = self.example_user('othello')
        do_add_alert_words(user2, ['another'])

        realm_words = alert_words_in_realm(user2.realm)
        self.assertEqual(len(realm_words), 2)
        self.assertEqual(set(realm_words.keys()), {user1.id, user2.id})
        self.assertEqual(set(realm_words[user1.id]),
                         set(self.interesting_alert_word_list))
        self.assertEqual(set(realm_words[user2.id]), {'another'})

    def test_json_list_default(self) -> None:
        self.login('hamlet')

        result = self.client_get('/json/users/me/alert_words')
        self.assert_json_success(result)
        self.assertEqual(result.json()['alert_words'], [])

    def test_json_list_nonempty(self) -> None:
        hamlet = self.example_user('hamlet')
        do_add_alert_words(hamlet, ['one', 'two', 'three'])

        self.login('hamlet')
        result = self.client_get('/json/users/me/alert_words')
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

    def test_json_list_add(self) -> None:
        self.login('hamlet')

        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['one ', '\n two', 'three'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

    def test_json_list_remove(self) -> None:
        self.login('hamlet')

        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['one', 'two', 'three'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

        result = self.client_delete('/json/users/me/alert_words',
                                    {'alert_words': ujson.dumps(['one'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'two', 'three'})

    def message_does_alert(self, user_profile: UserProfile, message: str) -> bool:
        """Send a bunch of messages as othello, so Hamlet is notified"""
        self.send_stream_message(self.example_user("othello"), "Denmark", message)
        user_message = most_recent_usermessage(user_profile)
        return 'has_alert_word' in user_message.flags_list()

    def test_alert_flags(self) -> None:
        self.login('hamlet')
        user_profile_hamlet = self.example_user('hamlet')

        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['one', 'two', 'three'])})
        self.assert_json_success(result)
        self.assertEqual(set(result.json()['alert_words']), {'one', 'two', 'three'})

        # Alerts in the middle of messages work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Normal alert one time"))
        # Alerts at the end of messages work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Normal alert one"))
        # Alerts at the beginning of messages work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "two normal alerts"))
        # Alerts with surrounding punctuation work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "This one? should alert"))
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Definitely time for three."))
        # Multiple alerts in a message work.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "One two three o'clock"))
        # Alerts are case-insensitive.
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "One o'clock"))
        self.assertTrue(self.message_does_alert(user_profile_hamlet, "Case of ONE, won't stop me"))

        # We don't cause alerts for matches in URLs.
        self.assertFalse(self.message_does_alert(user_profile_hamlet, "Don't alert on http://t.co/one/ urls"))
        self.assertFalse(self.message_does_alert(user_profile_hamlet, "Don't alert on http://t.co/one urls"))

    def test_update_alert_words(self) -> None:
        user_profile = self.example_user('hamlet')

        self.login_user(user_profile)
        result = self.client_post('/json/users/me/alert_words',
                                  {'alert_words': ujson.dumps(['ALERT'])})

        content = 'this is an ALERT for you'
        self.send_stream_message(user_profile, "Denmark", content)
        self.assert_json_success(result)

        original_message = most_recent_message(user_profile)

        user_message = most_recent_usermessage(user_profile)
        self.assertIn('has_alert_word', user_message.flags_list())

        result = self.client_patch("/json/messages/" + str(original_message.id), {
            'message_id': original_message.id,
            'content': 'new ALERT for you',
        })
        self.assert_json_success(result)

        user_message = most_recent_usermessage(user_profile)
        self.assertEqual(user_message.message.content, 'new ALERT for you')
        self.assertIn('has_alert_word', user_message.flags_list())

        result = self.client_patch("/json/messages/" + str(original_message.id), {
            'message_id': original_message.id,
            'content': 'sorry false alarm',
        })
        self.assert_json_success(result)

        user_message = most_recent_usermessage(user_profile)
        self.assertEqual(user_message.message.content, 'sorry false alarm')
        self.assertNotIn('has_alert_word', user_message.flags_list())
<|endoftext|>