<commit_before>import io
# import os
import logging
import xml.etree.cElementTree as ET
import requests
import zipfile
from ciso8601 import parse_datetime
from django.core.management.base import BaseCommand
from django.db.utils import IntegrityError
# from bustimes.utils import write_file
from ... import models

logger = logging.getLogger(__name__)

def get_atco_code(stop):
    ref = stop["@ref"]
    assert ref.startswith("atco:")
    return ref[5:]

def get_user_profile(element):
    return models.UserProfile.objects.get_or_create(
        {
            'name': element.findtext('Name'),
            'min_age': element.findtext('MinimumAge'),
            'max_age': element.findtext('MaximumAge'),
        },
        code=element.attrib["id"]
    )

class Command(BaseCommand):
    base_url = "https://data.bus-data.dft.gov.uk"

    def handle_file(self, source, open_file, filename=None):
        iterator = ET.iterparse(open_file)

        if not filename:
            filename = open_file.name

        for _, element in iterator:
            if element.tag[:31] == '{http://www.netex.org.uk/netex}':
                element.tag = element.tag[31:]

        user_profiles = {}
        for usage_parameter in element.find("dataObjects/CompositeFrame/frames/FareFrame/usageParameters"):
            user_profile, created = get_user_profile(usage_parameter)
            user_profiles[user_profile.code] = user_profile

        sales_offer_packages = {}
        for sales_offer_package in element.find("dataObjects/CompositeFrame/frames/FareFrame/salesOfferPackages"):
            sales_offer_package, created = models.SalesOfferPackage.objects.get_or_create(
                code=sales_offer_package.attrib["id"],
                name=sales_offer_package.findtext("Name", ""),
                description=sales_offer_package.findtext("Description", "")
            )
            sales_offer_packages[sales_offer_package.code] = sales_offer_package

        price_groups = {}
        price_group_prices = {}
        price_groups_element = element.find("dataObjects/CompositeFrame/frames/FareFrame/priceGroups")
        if price_groups_element:
            for price_group_element in price_groups_element:
                price_element = price_group_element.find("members/GeographicalIntervalPrice")  # assume only 1 ~
                price_group, created = models.PriceGroup.objects.get_or_create(
                    code=price_group_element.attrib["id"],
                    amount=price_element.findtext("Amount")
                )
                price_groups[price_group.code] = price_group
                price_group_prices[price_element.attrib["id"]] = price_group

        fare_zones = {}
        fare_zones_element = element.find("dataObjects/CompositeFrame/frames/FareFrame/fareZones")
        if fare_zones_element:
            for fare_zone in fare_zones_element:
                fare_zone, created = models.FareZone.objects.get_or_create(
                    code=fare_zone.attrib['id'], name=fare_zone.findtext("Name")
                )
                fare_zones[fare_zone.code] = fare_zone

        tariffs = {}
        distance_matrix_elements = {}
        for tariff_element in element.find("dataObjects/CompositeFrame/frames/FareFrame/tariffs"):
            fare_structre_elements = tariff_element.find("fareStructureElements")

            user_profile = fare_structre_elements.find(
                "FareStructureElement/GenericParameterAssignment/limitations/UserProfile"
            )
            user_profile, created = get_user_profile(user_profile)
            user_profiles[user_profile.code] = user_profile

            round_trip = fare_structre_elements.find(
                "FareStructureElement/GenericParameterAssignment/limitations/RoundTrip"
            )
            if round_trip:
                trip_type = round_trip.findtext('TripType')
            else:
                trip_type = ''

            tariff = models.Tariff.objects.create(
                code=tariff_element.attrib['id'], name=tariff_element.findtext("Name"),
                source=source, filename=filename,
                trip_type=trip_type, user_profile=user_profile
            )

            distance_matrix_element_elements = fare_structre_elements.find(
                "FareStructureElement/distanceMatrixElements"
            )
            if distance_matrix_element_elements:
                for distance_matrix_element in distance_matrix_element_elements:
                    start_zone = distance_matrix_element.find("StartTariffZoneRef").attrib["ref"]
                    end_zone = distance_matrix_element.find("EndTariffZoneRef").attrib["ref"]
                    price_group = distance_matrix_element.find("priceGroups/PriceGroupRef").attrib['ref']
                    distance_matrix_element, created = models.DistanceMatrixElement.objects.get_or_create(
                        code=distance_matrix_element.attrib["id"],
                        start_zone=fare_zones[start_zone],
                        end_zone=fare_zones[end_zone],
                        price_group=price_groups[price_group],
                        tariff=tariff,
                    )
                    distance_matrix_elements[distance_matrix_element.code] = distance_matrix_element

            tariffs[tariff.code] = tariff

        for fare_table_element in element.find("dataObjects/CompositeFrame/frames/FareFrame/fareTables"):
            tariff_ref = fare_table_element.find("usedIn/TariffRef")
            if tariff_ref is None:
                continue
            tariff_ref = tariff_ref.attrib["ref"]
            user_profile_ref = fare_table_element.find("pricesFor/UserProfileRef").attrib["ref"]
            if user_profile_ref not in user_profiles:
                user_profile_ref = f"fxc:{user_profile_ref}"
            sales_offer_package_ref = fare_table_element.find("pricesFor/SalesOfferPackageRef").attrib["ref"]
            table, created = models.FareTable.objects.update_or_create(
                {
                    "user_profile": user_profiles[user_profile_ref],
                    "sales_offer_package": sales_offer_packages[sales_offer_package_ref],
                    "description": fare_table_element.findtext("Description", "")
                },
                tariff=tariffs[tariff_ref],
                code=fare_table_element.attrib["id"],
                name=fare_table_element.findtext("Name", "")
            )

            if not created:
                table.column_set.all().delete()
                table.row_set.all().delete()

            columns = {}
            for column in fare_table_element.find('columns'):
                column = models.Column.objects.create(
                    table=table,
                    code=column.attrib['id'],
                    name=column.findtext('Name'),
                    order=column.attrib['order']
                )
                columns[column.code] = column

            rows = {}
            for row in fare_table_element.find('rows'):
                row = models.Row.objects.create(
                    table=table,
                    code=row.attrib['id'],
                    name=row.findtext('Name'),
                    order=row.attrib['order']
                )
                rows[row.code] = row

            for table in fare_table_element.find('includes'):
                cells = table.find('cells')
                if not cells:
                    continue
                for cell_element in cells:
                    columnn_ref = cell_element.find('ColumnRef').attrib['ref']
                    row_ref = cell_element.find('RowRef').attrib['ref']
                    distance_matrix_element_price = cell_element.find("DistanceMatrixElementPrice")
                    price_ref = distance_matrix_element_price.find("GeographicalIntervalPriceRef")
                    distance_matrix_element_ref = distance_matrix_element_price.find("DistanceMatrixElementRef")

                    if row_ref in rows:
                        row = rows[row_ref]
                    else:
                        row_ref_suffix = row_ref.split('@')[-1]
                        row_ref_suffix = f'@{row_ref_suffix}'
                        row_refs = [row_ref for row_ref in rows if row_ref.endswith(row_ref_suffix)]
                        assert len(row_refs) == 1
                        row = rows[row_refs[0]]
                    models.Cell.objects.create(
                        column=columns[columnn_ref],
                        row=row,
                        price_group=price_group_prices[price_ref.attrib["ref"]],
                        distance_matrix_element=distance_matrix_elements[distance_matrix_element_ref.attrib["ref"]]
                    )

    @staticmethod
    def add_arguments(parser):
        parser.add_argument('api_key', type=str)

    def handle_dataset(self, item):
        dataset_url = f"{self.base_url}/fares/dataset/{item['id']}/"
        modified = parse_datetime(item["modified"])
        try:
            dataset = models.DataSet.objects.get(url=dataset_url)
            if dataset.datetime == modified:
                return
            # has changed
            dataset.tariff_set.all().delete()
        except models.DataSet.DoesNotExist:
            dataset = models.DataSet(name=item["name"], description=item["description"], url=dataset_url)
            dataset.save()
        print(dataset_url)

        try:
            dataset.operators.set(item["noc"])
        except IntegrityError:
            print(item["noc"])

        download_url = f"{dataset_url}download/"
        response = self.session.get(download_url, stream=True)

        if response.headers['Content-Type'] == 'text/xml':
            # maybe not fully RFC 6266 compliant
            filename = response.headers['Content-Disposition'].split('filename', 1)[1][2:-1]
            print(filename)
            try:
                self.handle_file(dataset, response.raw, filename)
            except AttributeError as e:
                logger.error(e, exc_info=True)
                return
        else:
            assert response.headers['Content-Type'] == 'application/zip'
            with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
                for filename in archive.namelist():
                    print(' ', filename)
                    try:
                        self.handle_file(dataset, archive.open(filename), filename)
                    except AttributeError as e:
                        logger.error(e, exc_info=True)
                        return

        dataset.datetime = modified
        dataset.save(update_fields=["datetime"])

    def handle(self, api_key=None, **kwargs):
        # url = 'https://data.bus-data.dft.gov.uk/fares/download/bulk_archive'

        # with requests.get(url, stream=True) as response:
        #     # maybe not fully RFC 6266 compliant
        #     filename = response.headers['Content-Disposition'].split('filename', 1)[1][2:-1]

        #     if not os.path.exists(filename):
        #         write_file(filename, response)

        # print(filename)

        # with zipfile.ZipFile(filename) as archive:
        #     for filename in archive.namelist():
        #         print(' ', filename)
        #         # try:
        #         #     self.handle_file(dataset, archive.open(filename))
        #         # except AttributeError as e:
        #         #     logger.error(e, exc_info=True)

        # return

        self.session = requests.Session()

        url = f"{self.base_url}/api/v1/fares/dataset/"
        params = {
            'api_key': api_key,
            'status': 'published'
        }
        while url:
            print(url)
            response = self.session.get(url, params=params)

            data = response.json()

            for item in data['results']:
                self.handle_dataset(item)

            url = data['next']
            params = None
<commit_msg>I think I was trying to support more ticket types, can't remember<commit_after>import io
# import os
import logging
import xml.etree.cElementTree as ET
import requests
import zipfile
from ciso8601 import parse_datetime
from django.core.management.base import BaseCommand
from django.db.utils import IntegrityError
# from bustimes.utils import write_file
from ... import models

logger = logging.getLogger(__name__)

def get_atco_code(stop):
    ref = stop["@ref"]
    assert ref.startswith("atco:")
    return ref[5:]

def get_user_profile(element):
    return models.UserProfile.objects.get_or_create(
        {
            'name': element.findtext('Name'),
            'min_age': element.findtext('MinimumAge'),
            'max_age': element.findtext('MaximumAge'),
        },
        code=element.attrib["id"]
    )

class Command(BaseCommand):
    base_url = "https://data.bus-data.dft.gov.uk"

    def handle_file(self, source, open_file, filename=None):
        iterator = ET.iterparse(open_file)

        if not filename:
            filename = open_file.name

        for _, element in iterator:
            if element.tag[:31] == '{http://www.netex.org.uk/netex}':
                element.tag = element.tag[31:]

        user_profiles = {}
        for usage_parameter in element.find("dataObjects/CompositeFrame/frames/FareFrame/usageParameters"):
            user_profile, created = get_user_profile(usage_parameter)
            user_profiles[user_profile.code] = user_profile

        sales_offer_packages = {}
        for sales_offer_package in element.find("dataObjects/CompositeFrame/frames/FareFrame/salesOfferPackages"):
            sales_offer_package, created = models.SalesOfferPackage.objects.get_or_create(
                code=sales_offer_package.attrib["id"],
                name=sales_offer_package.findtext("Name", ""),
                description=sales_offer_package.findtext("Description", "")
            )
            sales_offer_packages[sales_offer_package.code] = sales_offer_package

        price_groups = {}
        price_group_prices = {}
        price_groups_element = element.find("dataObjects/CompositeFrame/frames/FareFrame/priceGroups")
        if price_groups_element:
            for price_group_element in price_groups_element:
                price_element = price_group_element.find("members/GeographicalIntervalPrice")  # assume only 1 ~
                price_group, created = models.PriceGroup.objects.get_or_create(
                    code=price_group_element.attrib["id"],
                    amount=price_element.findtext("Amount")
                )
                price_groups[price_group.code] = price_group
                price_group_prices[price_element.attrib["id"]] = price_group

        fare_zones = {}
        fare_zones_element = element.find("dataObjects/CompositeFrame/frames/FareFrame/fareZones")
        if fare_zones_element:
            for fare_zone in fare_zones_element:
                fare_zone, created = models.FareZone.objects.get_or_create(
                    code=fare_zone.attrib['id'], name=fare_zone.findtext("Name")
                )
                fare_zones[fare_zone.code] = fare_zone

        tariffs = {}
        distance_matrix_elements = {}
        for tariff_element in element.find("dataObjects/CompositeFrame/frames/FareFrame/tariffs"):
            fare_structre_elements = tariff_element.find("fareStructureElements")

            user_profile = fare_structre_elements.find(
                "FareStructureElement/GenericParameterAssignment/limitations/UserProfile"
            )
            user_profile, created = get_user_profile(user_profile)
            user_profiles[user_profile.code] = user_profile

            round_trip = fare_structre_elements.find(
                "FareStructureElement/GenericParameterAssignment/limitations/RoundTrip"
            )
            if round_trip:
                trip_type = round_trip.findtext('TripType')
            else:
                trip_type = ''

            tariff = models.Tariff.objects.create(
                code=tariff_element.attrib['id'], name=tariff_element.findtext("Name"),
                source=source, filename=filename,
                trip_type=trip_type, user_profile=user_profile
            )

            distance_matrix_element_elements = fare_structre_elements.find(
                "FareStructureElement/distanceMatrixElements"
            )
            if distance_matrix_element_elements:
                for distance_matrix_element in distance_matrix_element_elements:
                    start_zone = distance_matrix_element.find("StartTariffZoneRef").attrib["ref"]
                    end_zone = distance_matrix_element.find("EndTariffZoneRef").attrib["ref"]
                    price_group = distance_matrix_element.find("priceGroups/PriceGroupRef").attrib['ref']
                    distance_matrix_element, created = models.DistanceMatrixElement.objects.get_or_create(
                        code=distance_matrix_element.attrib["id"],
                        start_zone=fare_zones[start_zone],
                        end_zone=fare_zones[end_zone],
                        price_group=price_groups[price_group],
                        tariff=tariff,
                    )
                    distance_matrix_elements[distance_matrix_element.code] = distance_matrix_element

            tariffs[tariff.code] = tariff

        for fare_table_element in element.find("dataObjects/CompositeFrame/frames/FareFrame/fareTables"):
            tariff_ref = fare_table_element.find("usedIn/TariffRef")
            if tariff_ref is None:
                pass
            else:
                tariff_ref = tariff_ref.attrib["ref"]
                tariff = tariffs[tariff_ref]

            user_profile_ref = fare_table_element.find("pricesFor/UserProfileRef")
            if user_profile_ref is not None:
                user_profile_ref = user_profile_ref.attrib["ref"]
                if user_profile_ref not in user_profiles:
                    user_profile_ref = f"fxc:{user_profile}"
                user_profile = user_profiles[user_profile_ref]
            else:
                user_profile = None

            sales_offer_package_ref = fare_table_element.find("pricesFor/SalesOfferPackageRef")
            if sales_offer_package_ref is not None:
                sales_offer_package_ref = sales_offer_package_ref.attrib["ref"]
                sales_offer_package = sales_offer_packages[sales_offer_package_ref]
            else:
                sales_offer_package = None

            table, created = models.FareTable.objects.update_or_create(
                {
                    "user_profile": user_profile,
                    "sales_offer_package": sales_offer_package,
                    "description": fare_table_element.findtext("Description", "")
                },
                tariff=tariff,
                code=fare_table_element.attrib["id"],
                name=fare_table_element.findtext("Name", "")
            )

            if not created:
                table.column_set.all().delete()
                table.row_set.all().delete()

            columns = {}
            columns_element = fare_table_element.find('columns')
            if columns_element:
                for column in columns_element:
                    column = models.Column.objects.create(
                        table=table,
                        code=column.attrib['id'],
                        name=column.findtext('Name'),
                        order=column.attrib.get('order')
                    )
                    columns[column.code] = column

            rows = {}
            rows_element = fare_table_element.find('rows')
            if rows_element:
                for row in rows_element:
                    row = models.Row.objects.create(
                        table=table,
                        code=row.attrib['id'],
                        name=row.findtext('Name'),
                        order=row.attrib.get('order')
                    )
                    rows[row.code] = row

            for table in fare_table_element.find('includes'):
                cells_element = table.find('cells')
                if not cells_element:
                    continue
                for cell_element in cells_element:
                    columnn_ref = cell_element.find('ColumnRef').attrib['ref']
                    row_ref = cell_element.find('RowRef').attrib['ref']
                    distance_matrix_element_price = cell_element.find("DistanceMatrixElementPrice")
                    price_ref = distance_matrix_element_price.find("GeographicalIntervalPriceRef")
                    distance_matrix_element_ref = distance_matrix_element_price.find("DistanceMatrixElementRef")

                    if row_ref in rows:
                        row = rows[row_ref]
                    else:
                        # sometimes the RowRef doesn't correspond exactly to a Row id
                        row_ref_suffix = row_ref.split('@')[-1]
                        row_ref_suffix = f'@{row_ref_suffix}'
                        row_refs = [row_ref for row_ref in rows if row_ref.endswith(row_ref_suffix)]
                        assert len(row_refs) == 1
                        row = rows[row_refs[0]]
                    models.Cell.objects.create(
                        column=columns[columnn_ref],
                        row=row,
                        price_group=price_group_prices[price_ref.attrib["ref"]],
                        distance_matrix_element=distance_matrix_elements[distance_matrix_element_ref.attrib["ref"]]
                    )

    @staticmethod
    def add_arguments(parser):
        parser.add_argument('api_key', type=str)

    def handle_dataset(self, item):
        dataset_url = f"{self.base_url}/fares/dataset/{item['id']}/"
        modified = parse_datetime(item["modified"])
        try:
            dataset = models.DataSet.objects.get(url=dataset_url)
            if dataset.datetime == modified:
                return
            # has changed
            dataset.tariff_set.all().delete()
        except models.DataSet.DoesNotExist:
            dataset = models.DataSet(name=item["name"], description=item["description"], url=dataset_url)
            dataset.save()
        print(dataset_url)

        try:
            dataset.operators.set(item["noc"])
        except IntegrityError:
            print(item["noc"])

        download_url = f"{dataset_url}download/"
        response = self.session.get(download_url, stream=True)

        if response.headers['Content-Type'] == 'text/xml':
            # maybe not fully RFC 6266 compliant
            filename = response.headers['Content-Disposition'].split('filename', 1)[1][2:-1]
            print(filename)
            try:
                self.handle_file(dataset, response.raw, filename)
            except AttributeError as e:
                logger.error(e, exc_info=True)
                return
        else:
            assert response.headers['Content-Type'] == 'application/zip'
            with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
                for filename in archive.namelist():
                    print(' ', filename)
                    try:
                        self.handle_file(dataset, archive.open(filename), filename)
                    except AttributeError as e:
                        logger.error(e, exc_info=True)
                        return

        dataset.datetime = modified
        dataset.save(update_fields=["datetime"])

    def handle(self, api_key=None, **kwargs):
        # url = 'https://data.bus-data.dft.gov.uk/fares/download/bulk_archive'

        # with requests.get(url, stream=True) as response:
        #     # maybe not fully RFC 6266 compliant
        #     filename = response.headers['Content-Disposition'].split('filename', 1)[1][2:-1]

        #     if not os.path.exists(filename):
        #         write_file(filename, response)

        # print(filename)

        # with zipfile.ZipFile(filename) as archive:
        #     for filename in archive.namelist():
        #         print(' ', filename)
        #         # try:
        #         #     self.handle_file(dataset, archive.open(filename))
        #         # except AttributeError as e:
        #         #     logger.error(e, exc_info=True)

        # return

        self.session = requests.Session()

        url = f"{self.base_url}/api/v1/fares/dataset/"
        params = {
            'api_key': api_key,
            'status': 'published'
        }
        while url:
            print(url)
            response = self.session.get(url, params=params)

            data = response.json()

            for item in data['results']:
                self.handle_dataset(item)

            url = data['next']
            params = None
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding:utf8
import sys
reload(sys)
sys.setdefaultencoding( "utf8" )

import itchat
from itchat.content import *

import logging
import logging.handlers
from subprocess import call
formatter = logging.Formatter('%(asctime)s  %(message)s')

class MyTimedRotatingFileHandler(logging.handlers.TimedRotatingFileHandler):
    def doRollover(self):
        super(MyTimedRotatingFileHandler,self).doRollover()
        main_logger.info("Records log has been rotating...")
        call('git add .', shell = True)
        call('git commit -m "commiting..."', shell = True)
        call('git push origin master', shell = True)

def setup_logger(name, log_file):
    """Function setup as many loggers as you want"""

    handler = MyTimedRotatingFileHandler(log_file, when='midnight')
    handler.suffix = "%Y-%m-%d.log"
    handler.setFormatter(formatter)

    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)
    logger.addHandler(handler)

    return logger

# 自动回复文本等类别的群聊消息
# isGroupChat=True表示为群聊消息
@itchat.msg_register([TEXT, SHARING], isGroupChat=True)
def group_reply_text(msg):
    # 消息来自于哪个群聊
    chatroom_id = msg['FromUserName']
    # 发送者的昵称
    username = msg['ActualNickName']

    #from_chatroom = ''

    # 消息并不是来自于需要同步的群
    if not chatroom_id in chatroom_ids:
        return

    #for item in chatrooms:
    #    if item['UserName'] == chatroom_id:
    #        from_chatroom = item['NickName']

    # 根据消息类型转发至其他需要同步消息的群聊
    if msg['Type'] == TEXT:
        for item in chatrooms:
            if not item['UserName'] == chatroom_id:
                #itchat.send('%s from %s 说: \n%s' % (username, from_chatroom, msg['Content']), item['UserName'])
                itchat.send('%s 说: \n%s' % (username, msg['Content']), item['UserName'])
                #record_logger.info('%s from %s 说: \n%s' % (username, from_chatroom, msg['Content']))
                record_logger.info('%s 说: \n%s' % (username, msg['Content']))

    elif msg['Type'] == SHARING:
        for item in chatrooms:
            if not item['UserName'] == chatroom_id:
                #itchat.send('%s from %s 分享链接:\n %s\n%s' % (username,from_chatroom, msg['Text'], msg['Url']), item['UserName'])
                #record_logger.info('%s from %s 分享链接:\n %s\n%s' % (username,from_chatroom, msg['Text'], msg['Url']))
                itchat.send('%s 分享链接:\n %s\n%s' % (username, msg['Text'], msg['Url']), item['UserName'])
                record_logger.info('%s 分享链接:\n %s\n%s' % (username, msg['Text'], msg['Url']))

# 自动回复图片等类别的群聊消息
# isGroupChat=True表示为群聊消息
@itchat.msg_register([PICTURE, ATTACHMENT, VIDEO], isGroupChat=True)
def group_reply_media(msg):
    # 消息来自于哪个群聊
    chatroom_id = msg['FromUserName']
    # 发送者的昵称
    username = msg['ActualNickName']

    from_chatroom = ''

    for item in chatrooms:
        if item['UserName'] == chatroom_id:
            from_chatroom = item['NickName']

    # 消息并不是来自于需要同步的群
    if not chatroom_id in chatroom_ids:
        return

    # 如果为gif图片则不转发
    if msg['FileName'][-4:] == '.gif':
        return

    # 下载图片等文件
    msg['Text'](msg['FileName'])
    # 转发至其他需要同步消息的群聊
    for item in chatrooms:
        if not item['UserName'] == chatroom_id:
            #itchat.send('%s from %s 发送了：\n' % (username, from_chatroom), item['UserName'])
            #itchat.send('@%s@%s' % ({'Picture': 'img', 'Video': 'vid'}.get(msg['Type'], 'fil'), msg['FileName']), item['UserName'])
            itchat.send('%s 发送了：\n' % (username), item['UserName'])
            itchat.send('@%s@%s' % ({'Picture': 'img', 'Video': 'vid'}.get(msg['Type'], 'fil'), msg['FileName']), item['UserName'])

if __name__ == '__main__':
    #### main func ###
    main_logger = setup_logger("main_logger","main")
    record_logger = setup_logger('group_logger', 'redisgroup')
    # 扫二维码登录
    itchat.auto_login(hotReload=True,enableCmdQR=2)

    # 获取所有通讯录中的相关群聊
    chatrooms = itchat.search_chatrooms(name="Redis技术交流群")
    chatroom_ids = [c['UserName'] for c in chatrooms]

    main_logger.info('正在监测的群聊：%d 个' %(len(chatrooms)))
    main_logger.info(' '.join([item['NickName'] for item in chatrooms]))

    # 开始监测
    itchat.run()
<commit_msg>exclude pkl when pushing<commit_after>#!/usr/bin/env python
# coding:utf8
import sys
reload(sys)
sys.setdefaultencoding( "utf8" )

import itchat
from itchat.content import *

import logging
import logging.handlers
from subprocess import call
formatter = logging.Formatter('%(asctime)s  %(message)s')

class MyTimedRotatingFileHandler(logging.handlers.TimedRotatingFileHandler):
    def doRollover(self):
        super(MyTimedRotatingFileHandler,self).doRollover()
        main_logger.info("Records log has been rotating...")
        call('git add -u', shell = True)
        call('git reset -- ChatRecords/itchat.pkl', shell = True)
        call('git commit -m "commiting..."', shell = True)
        call('git push origin master', shell = True)

def setup_logger(name, log_file):
    """Function setup as many loggers as you want"""

    handler = MyTimedRotatingFileHandler(log_file, when='midnight')
    handler.suffix = "%Y-%m-%d.log"
    handler.setFormatter(formatter)

    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)
    logger.addHandler(handler)

    return logger

# 自动回复文本等类别的群聊消息
# isGroupChat=True表示为群聊消息
@itchat.msg_register([TEXT, SHARING], isGroupChat=True)
def group_reply_text(msg):
    # 消息来自于哪个群聊
    chatroom_id = msg['FromUserName']
    # 发送者的昵称
    username = msg['ActualNickName']

    #from_chatroom = ''

    # 消息并不是来自于需要同步的群
    if not chatroom_id in chatroom_ids:
        return

    #for item in chatrooms:
    #    if item['UserName'] == chatroom_id:
    #        from_chatroom = item['NickName']

    # 根据消息类型转发至其他需要同步消息的群聊
    if msg['Type'] == TEXT:
        for item in chatrooms:
            if not item['UserName'] == chatroom_id:
                #itchat.send('%s from %s 说: \n%s' % (username, from_chatroom, msg['Content']), item['UserName'])
                itchat.send('%s 说: \n%s' % (username, msg['Content']), item['UserName'])
                #record_logger.info('%s from %s 说: \n%s' % (username, from_chatroom, msg['Content']))
                record_logger.info('%s 说: \n%s' % (username, msg['Content']))

    elif msg['Type'] == SHARING:
        for item in chatrooms:
            if not item['UserName'] == chatroom_id:
                #itchat.send('%s from %s 分享链接:\n %s\n%s' % (username,from_chatroom, msg['Text'], msg['Url']), item['UserName'])
                #record_logger.info('%s from %s 分享链接:\n %s\n%s' % (username,from_chatroom, msg['Text'], msg['Url']))
                itchat.send('%s 分享链接:\n %s\n%s' % (username, msg['Text'], msg['Url']), item['UserName'])
                record_logger.info('%s 分享链接:\n %s\n%s' % (username, msg['Text'], msg['Url']))

# 自动回复图片等类别的群聊消息
# isGroupChat=True表示为群聊消息
@itchat.msg_register([PICTURE, ATTACHMENT, VIDEO], isGroupChat=True)
def group_reply_media(msg):
    # 消息来自于哪个群聊
    chatroom_id = msg['FromUserName']
    # 发送者的昵称
    username = msg['ActualNickName']

    from_chatroom = ''

    for item in chatrooms:
        if item['UserName'] == chatroom_id:
            from_chatroom = item['NickName']

    # 消息并不是来自于需要同步的群
    if not chatroom_id in chatroom_ids:
        return

    # 如果为gif图片则不转发
    if msg['FileName'][-4:] == '.gif':
        return

    # 下载图片等文件
    msg['Text'](msg['FileName'])
    # 转发至其他需要同步消息的群聊
    for item in chatrooms:
        if not item['UserName'] == chatroom_id:
            #itchat.send('%s from %s 发送了：\n' % (username, from_chatroom), item['UserName'])
            #itchat.send('@%s@%s' % ({'Picture': 'img', 'Video': 'vid'}.get(msg['Type'], 'fil'), msg['FileName']), item['UserName'])
            itchat.send('%s 发送了：\n' % (username), item['UserName'])
            itchat.send('@%s@%s' % ({'Picture': 'img', 'Video': 'vid'}.get(msg['Type'], 'fil'), msg['FileName']), item['UserName'])

if __name__ == '__main__':
    #### main func ###
    main_logger = setup_logger("main_logger","main")
    record_logger = setup_logger('group_logger', 'redisgroup')
    # 扫二维码登录
    itchat.auto_login(hotReload=True,enableCmdQR=2)

    # 获取所有通讯录中的相关群聊
    chatrooms = itchat.search_chatrooms(name="Redis技术交流群")
    chatroom_ids = [c['UserName'] for c in chatrooms]

    main_logger.info('正在监测的群聊：%d 个' %(len(chatrooms)))
    main_logger.info(' '.join([item['NickName'] for item in chatrooms]))

    # 开始监测
    itchat.run()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
    vdirsyncer.utils.vobject
    ~~~~~~~~~~~~~~~~~~~~~~~~

    :copyright: (c) 2014 Markus Unterwaditzer & contributors
    :license: MIT, see LICENSE for more details.
'''
import hashlib

import icalendar.cal
import icalendar.caselessdict
import icalendar.parser

from . import cached_property, split_sequence
from .compat import itervalues, text_type

def _process_properties(*s):
    rv = set()
    for key in s:
        rv.add(key + ':')
        rv.add(key + ';')

    return tuple(rv)

IGNORE_PROPS = _process_properties(
    # PRODID is changed by radicale for some reason after upload
    'PRODID',
    # VERSION can get lost in singlefile storage
    'VERSION',
    # X-RADICALE-NAME is used by radicale, because hrefs don't really exist in
    # their filesystem backend
    'X-RADICALE-NAME',
    # REV is from the VCARD specification and is supposed to change when the
    # item does -- however, we can determine that ourselves
    'REV'
)

# Whether the installed icalendar version has
# https://github.com/collective/icalendar/pull/136
# (support for keeping the order of properties and parameters)
#
# This basically checks whether the superclass of all icalendar classes has a
# method from OrderedDict.
ICALENDAR_ORIGINAL_ORDER_SUPPORT = \
    hasattr(icalendar.caselessdict.CaselessDict, '__reversed__')

class Item(object):

    '''Immutable wrapper class for VCALENDAR (VEVENT, VTODO) and
    VCARD'''

    def __init__(self, raw):
        assert isinstance(raw, text_type)
        self._raw = raw

    @cached_property
    def raw(self):
        '''Raw content of the item, as unicode string.

        Vdirsyncer doesn't validate the content in any way.
        '''
        return self._raw

    @cached_property
    def uid(self):
        '''Global identifier of the item, across storages, doesn't change after
        a modification of the item.'''
        stack = [self.parsed]
        while stack:
            component = stack.pop()
            if component is None:
                continue
            uid = component.get('UID', None)
            if uid:
                return uid
            stack.extend(component.subcomponents)

        for line in self.raw.splitlines():
            if line.startswith(u'UID:'):
                uid = line[4:].strip()
                if uid:
                    return uid

    @cached_property
    def hash(self):
        '''Hash of self.raw, used for etags.'''
        return hash_item(self.raw)

    @cached_property
    def ident(self):
        '''Used for generating hrefs and matching up items during
        synchronization. This is either the UID or the hash of the item's
        content.'''

        # We hash the item instead of directly using its raw content, because
        #
        # 1. The raw content might be really large, e.g. when its a contact
        #    with a picture, which bloats the status file.
        #
        # 2. The status file would contain really sensitive information.
        return self.uid or self.hash

    @cached_property
    def parsed(self):
        try:
            return icalendar.cal.Component.from_ical(self.raw)
        except Exception:
            return None

def normalize_item(item, ignore_props=IGNORE_PROPS, use_icalendar=True):
    if not isinstance(item, Item):
        item = Item(item)
    if use_icalendar and item.parsed is not None:
        # We have to explicitly check "is not None" here because VCALENDARS
        # with only subcomponents and no own properties are also false-ish.
        lines = to_unicode_lines(item.parsed)
    else:
        lines = sorted(item.raw.splitlines())

    return u'\r\n'.join(line.strip()
                        for line in lines
                        if line.strip() and
                        not line.startswith(IGNORE_PROPS))

def hash_item(text):
    return hashlib.sha256(normalize_item(text).encode('utf-8')).hexdigest()

def split_collection(text, inline=(u'VTIMEZONE',),
                     wrap_items_with=(u'VCALENDAR',)):
    '''Emits items in the order they occur in the text.'''
    assert isinstance(text, text_type)
    collections = icalendar.cal.Component.from_ical(text, multiple=True)
    collection_name = None

    for collection in collections:
        items = collection.subcomponents
        if collection_name is None:
            collection_name = collection.name

            if collection_name in wrap_items_with:
                start = u'BEGIN:{}'.format(collection_name)
                end = u'END:{}'.format(collection_name)
            else:
                start = end = u''

        if collection.name != collection_name:
            raise ValueError('Different types of components at top-level. '
                             'Expected {}, got {}.'
                             .format(collection_name, collection.name))

        inlined_items, normal_items = \
            split_sequence(items, lambda item: item.name in inline)

        for item in normal_items:
            lines = []
            lines.append(start)
            for inlined_item in inlined_items:
                lines.extend(to_unicode_lines(inlined_item))

            lines.extend(to_unicode_lines(item))
            lines.append(end)

            yield u''.join(line + u'\r\n' for line in lines if line)

def to_unicode_lines(item):
    '''icalendar doesn't provide an efficient way of getting the ical data as
    unicode. So let's do it ourselves.'''

    if ICALENDAR_ORIGINAL_ORDER_SUPPORT:
        content_lines = item.content_lines(sorted=False)
    else:
        content_lines = item.content_lines()

    for content_line in content_lines:
        if content_line:
            # https://github.com/untitaker/vdirsyncer/issues/70
            # XXX: icalendar escapes semicolons which are not supposed to get
            # escaped, because it is not aware of vcard
            content_line = content_line.replace(u'\\;', u';')
            yield icalendar.parser.foldline(content_line)

_default_join_wrappers = {
    u'VCALENDAR': (u'VCALENDAR', (u'VTIMEZONE',)),
    u'VEVENT': (u'VCALENDAR', (u'VTIMEZONE',)),
    u'VTODO': (u'VCALENDAR', (u'VTIMEZONE',)),
    u'VCARD': (u'VADDRESSBOOK', ())
}

def join_collection(items, wrappers=_default_join_wrappers):
    '''
    :param wrappers: {
        item_type: wrapper_type, common_components
    }

    Common components are those who can be moved from the item into the
    wrapper, with duplicates removed.
    '''
    inline = {}
    components = []
    wrapper_type = None
    inline_types = None
    item_type = None

    def handle_item(item):
        if item.name in inline_types:
            inline[tuple(to_unicode_lines(item))] = item
        else:
            components.append(item)

    for item in items:
        component = icalendar.cal.Component.from_ical(item)

        if item_type is None:
            item_type = component.name
            wrapper_type, inline_types = wrappers[item_type]

        if component.name == item_type:
            if item_type == wrapper_type:
                for subcomponent in component.subcomponents:
                    handle_item(subcomponent)
            else:
                handle_item(component)

    start = end = u''
    if wrapper_type is not None:
        start = u'BEGIN:{}'.format(wrapper_type)
        end = u'END:{}'.format(wrapper_type)

    lines = [start]
    for inlined_item in itervalues(inline):
        lines.extend(to_unicode_lines(inlined_item))
    for component in components:
        lines.extend(to_unicode_lines(component))
    lines.append(end)

    return u''.join(line + u'\r\n' for line in lines)
<commit_msg>More refactoring<commit_after># -*- coding: utf-8 -*-
'''
    vdirsyncer.utils.vobject
    ~~~~~~~~~~~~~~~~~~~~~~~~

    :copyright: (c) 2014 Markus Unterwaditzer & contributors
    :license: MIT, see LICENSE for more details.
'''
import hashlib
from itertools import chain, tee

import icalendar.cal
import icalendar.caselessdict
import icalendar.parser

from . import cached_property, split_sequence, uniq
from .compat import text_type

def _process_properties(*s):
    rv = set()
    for key in s:
        rv.add(key + ':')
        rv.add(key + ';')

    return tuple(rv)

IGNORE_PROPS = _process_properties(
    # PRODID is changed by radicale for some reason after upload
    'PRODID',
    # VERSION can get lost in singlefile storage
    'VERSION',
    # X-RADICALE-NAME is used by radicale, because hrefs don't really exist in
    # their filesystem backend
    'X-RADICALE-NAME',
    # REV is from the VCARD specification and is supposed to change when the
    # item does -- however, we can determine that ourselves
    'REV'
)

# Whether the installed icalendar version has
# https://github.com/collective/icalendar/pull/136
# (support for keeping the order of properties and parameters)
#
# This basically checks whether the superclass of all icalendar classes has a
# method from OrderedDict.
ICALENDAR_ORIGINAL_ORDER_SUPPORT = \
    hasattr(icalendar.caselessdict.CaselessDict, '__reversed__')

class Item(object):

    '''Immutable wrapper class for VCALENDAR (VEVENT, VTODO) and
    VCARD'''

    def __init__(self, raw):
        assert isinstance(raw, text_type)
        self._raw = raw

    @cached_property
    def raw(self):
        '''Raw content of the item, as unicode string.

        Vdirsyncer doesn't validate the content in any way.
        '''
        return self._raw

    @cached_property
    def uid(self):
        '''Global identifier of the item, across storages, doesn't change after
        a modification of the item.'''
        stack = [self.parsed]
        while stack:
            component = stack.pop()
            if component is None:
                continue
            uid = component.get('UID', None)
            if uid:
                return uid
            stack.extend(component.subcomponents)

        for line in self.raw.splitlines():
            if line.startswith(u'UID:'):
                uid = line[4:].strip()
                if uid:
                    return uid

    @cached_property
    def hash(self):
        '''Hash of self.raw, used for etags.'''
        return hash_item(self.raw)

    @cached_property
    def ident(self):
        '''Used for generating hrefs and matching up items during
        synchronization. This is either the UID or the hash of the item's
        content.'''

        # We hash the item instead of directly using its raw content, because
        #
        # 1. The raw content might be really large, e.g. when its a contact
        #    with a picture, which bloats the status file.
        #
        # 2. The status file would contain really sensitive information.
        return self.uid or self.hash

    @cached_property
    def parsed(self):
        try:
            return icalendar.cal.Component.from_ical(self.raw)
        except Exception:
            return None

def normalize_item(item, ignore_props=IGNORE_PROPS, use_icalendar=True):
    if not isinstance(item, Item):
        item = Item(item)
    if use_icalendar and item.parsed is not None:
        # We have to explicitly check "is not None" here because VCALENDARS
        # with only subcomponents and no own properties are also false-ish.
        lines = to_unicode_lines(item.parsed)
    else:
        lines = sorted(item.raw.splitlines())

    return u'\r\n'.join(line.strip()
                        for line in lines
                        if line.strip() and
                        not line.startswith(IGNORE_PROPS))

def hash_item(text):
    return hashlib.sha256(normalize_item(text).encode('utf-8')).hexdigest()

def split_collection(text, inline=(u'VTIMEZONE',),
                     wrap_items_with=(u'VCALENDAR',)):
    '''Emits items in the order they occur in the text.'''
    assert isinstance(text, text_type)
    collections = icalendar.cal.Component.from_ical(text, multiple=True)
    collection_name = None

    for collection in collections:
        if collection_name is None:
            collection_name = collection.name
            start = end = ()
            if collection_name in wrap_items_with:
                start = (u'BEGIN:{}'.format(collection_name),)
                end = (u'END:{}'.format(collection_name),)

        elif collection.name != collection_name:
            raise ValueError('Different types of components at top-level. '
                             'Expected {}, got {}.'
                             .format(collection_name, collection.name))

        inlined_items, normal_items = split_sequence(
            collection.subcomponents,
            lambda item: item.name in inline
        )
        inlined_lines = list(chain(*(to_unicode_lines(inlined_item)
                                     for inlined_item in inlined_items)))

        for item in normal_items:
            lines = chain(start, inlined_lines, to_unicode_lines(item), end)
            yield u''.join(line + u'\r\n' for line in lines if line)

def to_unicode_lines(item):
    '''icalendar doesn't provide an efficient way of getting the ical data as
    unicode. So let's do it ourselves.'''

    if ICALENDAR_ORIGINAL_ORDER_SUPPORT:
        content_lines = item.content_lines(sorted=False)
    else:
        content_lines = item.content_lines()

    for content_line in content_lines:
        if content_line:
            # https://github.com/untitaker/vdirsyncer/issues/70
            # XXX: icalendar escapes semicolons which are not supposed to get
            # escaped, because it is not aware of vcard
            content_line = content_line.replace(u'\\;', u';')
            yield icalendar.parser.foldline(content_line)

_default_join_wrappers = {
    u'VCALENDAR': u'VCALENDAR',
    u'VEVENT': u'VCALENDAR',
    u'VTODO': u'VCALENDAR',
    u'VCARD': u'VADDRESSBOOK'
}

def join_collection(items, wrappers=_default_join_wrappers):
    '''
    :param wrappers: {
        item_type: wrapper_type
    }
    '''

    items1, items2 = tee((icalendar.cal.Component.from_ical(x)
                          for x in items), 2)
    item_type, wrapper_type = _get_item_type(items1, wrappers)
    _get_item_components = lambda x: (x.name == wrapper_type
                                      and x.subcomponents
                                      or [x])
    components = chain(*(_get_item_components(x) for x in items2))
    lines = chain(*uniq(tuple(to_unicode_lines(x)) for x in components))

    if wrapper_type is not None:
        start = [u'BEGIN:{}'.format(wrapper_type)]
        end = [u'END:{}'.format(wrapper_type)]
        lines = chain(start, lines, end)
    return u''.join(line + u'\r\n' for line in lines)

def _get_item_type(components, wrappers):
    for component in components:
        try:
            item_type = component.name
            wrapper_type = wrappers[item_type]
        except KeyError:
            pass
        else:
            return item_type, wrapper_type
    return None, None
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Provides data sources to use for benchmarking."""

import ConfigParser
import glob
import os
import os.path
import random
import re
import subprocess
import sys
import time

# relative
import addr_util
import selectors
import util

# Pick the most accurate timer for a platform. Stolen from timeit.py:
if sys.platform[:3] == 'win':
  DEFAULT_TIMER = time.clock
else:
  DEFAULT_TIMER = time.time

GLOBAL_DATA_CACHE = {}

DEFAULT_CONFIG_PATH = 'config/data_sources.cfg'
MAX_NON_UNIQUE_RECORD_COUNT = 500000
MAX_FILE_MTIME_AGE_DAYS = 45
MIN_FILE_SIZE = 10000
MIN_RECOMMENDED_RECORD_COUNT = 200
MAX_FQDN_SYNTHESIZE_PERCENT = 4

class DataSources(object):
  """A collection of methods related to available hostname data sources."""

  def __init__(self, config_path=DEFAULT_CONFIG_PATH, status_callback=None):
    self.source_cache = GLOBAL_DATA_CACHE
    self.source_config = {}
    self.status_callback = status_callback
    self._LoadConfigFromPath(config_path)

  def msg(self, msg, **kwargs):
    if self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print '- %s' % msg

  def _LoadConfigFromPath(self, path):
    """Load a configuration file describing data sources that may be available."""
    conf_file = util.FindDataFile(path)
    config = ConfigParser.ConfigParser()
    config.read(conf_file)
    for section in config.sections():
      if section not in self.source_config:
        self.source_config[section] = {
            'name': None,
            'search_paths': set(),
            # Store whether or not this data source contains personal data
            'full_hostnames': True,
            'synthetic': False,
            'include_duplicates': False,
            'max_mtime_days': MAX_FILE_MTIME_AGE_DAYS
        }

      for (key, value) in config.items(section):
        if key == 'name':
          self.source_config[section]['name'] = value
        elif key == 'full_hostnames' and int(value) == 0:
          self.source_config[section]['full_hostnames'] = False
        elif key == 'max_mtime_days':
          self.source_config[section]['max_mtime_days'] = int(value)
        elif key == 'include_duplicates':
          self.source_config[section]['include_duplicates'] = bool(value)
        elif key == 'synthetic':
          self.source_config[section]['synthetic'] = bool(value)
        else:
          self.source_config[section]['search_paths'].add(value)

  def ListSourceTypes(self):
    """Get a list of all data sources we know about."""
    return sorted(self.source_config.keys())

  def ListSourcesWithDetails(self):
    """Get a list of all data sources found with total counts.

    Returns:
      List of tuples in form of (short_name, full_name, full_hosts, # of entries)
    """
    for source in self.ListSourceTypes():
      max_mtime = self.source_config[source]['max_mtime_days']
      self._GetHostsFromSource(source, min_file_size=MIN_FILE_SIZE,
                               max_mtime_age_days=max_mtime)

    details = []
    for source in self.source_cache:
      details.append((source,
                      self.source_config[source]['name'],
                      self.source_config[source]['synthetic'],
                      len(self.source_cache[source])))
    return sorted(details, key=lambda x: (x[2], x[3] * -1))

  def ListSourceTitles(self):
    """Return a list of sources in title + count format."""
    titles = []
    seen_synthetic = False
    seen_organic = False
    for (unused_type, name, is_synthetic, count) in self.ListSourcesWithDetails():
      if not is_synthetic:
        seen_organic = True

      if is_synthetic and seen_organic and not seen_synthetic:
        titles.append('-' * 36)
        seen_synthetic = True
      titles.append('%s (%s)' % (name, count))
    return titles

  def ConvertSourceTitleToType(self, detail):
    """Convert a detail name to a source type."""
    for source_type in self.source_config:
      if detail.startswith(self.source_config[source_type]['name']):
        return source_type

  def GetBestSourceDetails(self):
    return self.ListSourcesWithDetails()[0]

  def GetNameForSource(self, source):
    if source in self.source_config:
      return self.source_config[source]['name']
    else:
      # Most likely a custom file path
      return source

  def GetCachedRecordCountForSource(self, source):
    return len(self.source_cache[source])

  def _CreateRecordsFromHostEntries(self, entries, include_duplicates=False):
    """Create records from hosts, removing duplicate entries and IP's.

    Args:
      entries: A list of test-data entries.
      include_duplicates: Whether or not to filter duplicates (optional: False)

    Returns:
      A tuple of (filtered records, full_host_names (Boolean)

    Raises:
      ValueError: If no records could be grokked from the input.
    """
    last_entry = None

    records = []
    full_host_count = 0
    for entry in entries:
      if entry == last_entry and not include_duplicates:
        continue
      else:
        last_entry = entry

      if ' ' in entry:
        (record_type, host) = entry.split(' ')
      else:
        record_type = 'A'
        host = entry

      if not addr_util.IP_RE.match(host) and not addr_util.INTERNAL_RE.search(host):
        if not host.endswith('.'):
          host += '.'
        records.append((record_type, host))

        if addr_util.FQDN_RE.match(host):
          full_host_count += 1

    if not records:
      raise ValueError('No records could be created from: %s' % entries)

    # Now that we've read everything, are we dealing with domains or full hostnames?
    full_host_percent = full_host_count / float(len(records)) * 100
    if full_host_percent < MAX_FQDN_SYNTHESIZE_PERCENT:
      full_host_names = True
    else:
      full_host_names = False
    return (records, full_host_names)

  def GetTestsFromSource(self, source, count=50, select_mode=None):
    """Parse records from source, and return tuples to use for testing.

    Args:
      source: A source name (str) that has been configured.
      count: Number of tests to generate from the source (int)
      select_mode: automatic, weighted, random, chunk (str)

    Returns:
      A list of record tuples in the form of (req_type, hostname)

    Raises:
      ValueError: If no usable records are found from the data source.

    This is tricky because we support 3 types of input data:

    - List of domains
    - List of hosts
    - List of record_type + hosts
    """
    records = []

    if source in self.source_config:
      include_duplicates = self.source_config[source].get('include_duplicates', False)
    else:
      include_duplicates = False

    records = self._GetHostsFromSource(source)
    if not records:
      raise ValueError('Unable to generate records from %s (nothing found)' % source)

    self.msg('Generating tests from %s (%s records, selecting %s %s)'
             % (self.GetNameForSource(source), len(records), count, select_mode))
    (records, are_records_fqdn) = self._CreateRecordsFromHostEntries(records,
                                                                     include_duplicates=include_duplicates)
    # First try to resolve whether to use weighted or random.
    if select_mode in ('weighted', 'automatic', None):
      # If we are in include_duplicates mode (cachemiss, cachehit, etc.), we have different rules.
      if include_duplicates:
        if count > len(records):
          select_mode = 'random'
        else:
          select_mode = 'chunk'
      elif len(records) != len(set(records)):
        if select_mode == 'weighted':
          self.msg('%s data contains duplicates, switching select_mode to random' % source)
        select_mode = 'random'
      else:
        select_mode = 'weighted'

    self.msg('Selecting %s out of %s sanitized records (%s mode).' %
             (count, len(records), select_mode))
    if select_mode == 'weighted':
      records = selectors.WeightedDistribution(records, count)
    elif select_mode == 'chunk':
      records = selectors.ChunkSelect(records, count)
    elif select_mode == 'random':
      records = selectors.RandomSelect(records, count, include_duplicates=include_duplicates)
    else:
      raise ValueError('No such final selection mode: %s' % select_mode)

    # For custom filenames
    if source not in self.source_config:
      self.source_config[source] = {'synthetic': True}

    if are_records_fqdn:
      self.source_config[source]['full_hostnames'] = False
      self.msg('%s input appears to be predominantly domain names. Synthesizing FQDNs' % source)
      synthesized = []
      for (req_type, hostname) in records:
        if not addr_util.FQDN_RE.match(hostname):
          hostname = self._GenerateRandomHostname(hostname)
        synthesized.append((req_type, hostname))
      return synthesized
    else:
      return records

  def _GenerateRandomHostname(self, domain):
    """Generate a random hostname f or a given domain."""
    oracle = random.randint(0, 100)
    if oracle < 70:
      return 'www.%s' % domain
    elif oracle < 95:
      return domain
    elif oracle < 98:
      return 'static.%s' % domain
    else:
      return 'cache-%s.%s' % (random.randint(0, 10), domain)

  def _GetHostsFromSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Get data for a particular source. This needs to be fast.

    Args:
      source: A configured source type (str)
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      list of hostnames gathered from data source.

    The results of this function are cached by source type.
    """
    if source in self.source_cache:
      return self.source_cache[source]
    filename = self._FindBestFileForSource(source, min_file_size=min_file_size,
                                           max_mtime_age_days=max_mtime_age_days)
    if not filename:
      return None

    size_mb = os.path.getsize(filename) / 1024.0 / 1024.0
    self.msg('Reading %s: %s (%0.1fMB)' % (self.GetNameForSource(source), filename, size_mb))
    start_clock = DEFAULT_TIMER()
    if filename.endswith('.pcap') or filename.endswith('.tcp'):
      hosts = self._ExtractHostsFromPcapFile(filename)
    else:
      hosts = self._ExtractHostsFromHistoryFile(filename)

    if not hosts:
      hosts = self._ReadDataFile(filename)
    duration = DEFAULT_TIMER() - start_clock
    if duration > 5:
      self.msg('%s data took %1.1fs to read!' % (self.GetNameForSource(source), duration))
    self.source_cache[source] = hosts
    return hosts

  def _ExtractHostsFromHistoryFile(self, path):
    """Get a list of sanitized records from a history file containing URLs."""
    # This regexp is fairly general (no ip filtering), since we need speed more
    # than precision at this stage.
    parse_re = re.compile('https*://([\-\w]+\.[\-\w\.]+)')
    return parse_re.findall(open(path, 'rb').read())

  def _ExtractHostsFromPcapFile(self, path):
    """Get a list of requests out of a pcap file - requires tcpdump."""
    self.msg('Extracting requests from pcap file using tcpdump')
    cmd = 'tcpdump -r %s -n port 53' % path
    pipe = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).stdout
    parse_re = re.compile(' ([A-Z]+)\? ([\-\w\.]+)')
    requests = []
    for line in pipe:
      if '?' not in line:
        continue
      match = parse_re.search(line)
      if match:
        requests.append(' '.join(match.groups()))
    return requests

  def _ReadDataFile(self, path):
    """Read a line-based datafile."""
    records = []
    domains_re = re.compile('^\w[\w\.-]+[a-zA-Z]$')
    requests_re = re.compile('^[A-Z]{1,4} \w[\w\.-]+\.$')
    for line in open(path):
      if domains_re.match(line) or requests_re.match(line):
        records.append(line.rstrip())
    return records

  def _GetSourceSearchPaths(self, source):
    """Get a list of possible search paths (globs) for a given source."""

    # This is likely a custom file path
    if source not in self.source_config:
      return [source]

    search_paths = []
    environment_re = re.compile('%(\w+)%')

    # First get through resolving environment variables
    for path in self.source_config[source]['search_paths']:
      env_vars = set(environment_re.findall(path))
      if env_vars:
        for variable in env_vars:
          env_var = os.getenv(variable, False)
          if env_var:
            path = path.replace('%%%s%%' % variable, env_var)
          else:
            path = None

      # If everything is good, replace all '/'  chars with the os path variable.
      if path:
        path = path.replace('/', os.sep)
        search_paths.append(path)

        # This moment of weirdness brought to you by Windows XP(tm). If we find
        # a Local or Roaming keyword in path, add the other forms to the search
        # path.
        if sys.platform[:3] == 'win':
          keywords = ('Local', 'Roaming')
          for keyword in keywords:
            if keyword in path:
              replacement = keywords[keywords.index(keyword)-1]
              search_paths.append(path.replace('\\%s' % keyword, '\\%s' % replacement))
              search_paths.append(path.replace('\\%s' % keyword, ''))

    return search_paths

  def _FindBestFileForSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Find the best file (newest over X size) to use for a given source type.

    Args:
      source: source type
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      A file path.
    """
    found = []
    for path in self._GetSourceSearchPaths(source):
      if not os.path.isabs(path):
        path = util.FindDataFile(path)

      for filename in glob.glob(path):
        if min_file_size and os.path.getsize(filename) < min_file_size:
          self.msg('Skipping %s (only %sb)' % (filename, os.path.getsize(filename)))
        else:
          found.append(filename)

    if found:
      newest = sorted(found, key=os.path.getmtime)[-1]
      age_days = (time.time() - os.path.getmtime(newest)) / 86400
      if max_mtime_age_days and age_days > max_mtime_age_days:
        self.msg('Skipping %s (%2.0fd old)' % (newest, age_days))
      else:
        return newest
    else:
      return None

if __name__ == '__main__':
  parser = DataSources()
  print parser.ListSourceTypes()
  print parser.ListSourcesWithDetails()
  best = parser.ListSourcesWithDetails()[0][0]
  print len(parser.GetRecordsFromSource(best))
<commit_msg>Make keywords a list instead of a tuple. Fixes issue 127<commit_after>#!/usr/bin/env python
# Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Provides data sources to use for benchmarking."""

import ConfigParser
import glob
import os
import os.path
import random
import re
import subprocess
import sys
import time

# relative
import addr_util
import selectors
import util

# Pick the most accurate timer for a platform. Stolen from timeit.py:
if sys.platform[:3] == 'win':
  DEFAULT_TIMER = time.clock
else:
  DEFAULT_TIMER = time.time

GLOBAL_DATA_CACHE = {}

DEFAULT_CONFIG_PATH = 'config/data_sources.cfg'
MAX_NON_UNIQUE_RECORD_COUNT = 500000
MAX_FILE_MTIME_AGE_DAYS = 45
MIN_FILE_SIZE = 10000
MIN_RECOMMENDED_RECORD_COUNT = 200
MAX_FQDN_SYNTHESIZE_PERCENT = 4

class DataSources(object):
  """A collection of methods related to available hostname data sources."""

  def __init__(self, config_path=DEFAULT_CONFIG_PATH, status_callback=None):
    self.source_cache = GLOBAL_DATA_CACHE
    self.source_config = {}
    self.status_callback = status_callback
    self._LoadConfigFromPath(config_path)

  def msg(self, msg, **kwargs):
    if self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print '- %s' % msg

  def _LoadConfigFromPath(self, path):
    """Load a configuration file describing data sources that may be available."""
    conf_file = util.FindDataFile(path)
    config = ConfigParser.ConfigParser()
    config.read(conf_file)
    for section in config.sections():
      if section not in self.source_config:
        self.source_config[section] = {
            'name': None,
            'search_paths': set(),
            # Store whether or not this data source contains personal data
            'full_hostnames': True,
            'synthetic': False,
            'include_duplicates': False,
            'max_mtime_days': MAX_FILE_MTIME_AGE_DAYS
        }

      for (key, value) in config.items(section):
        if key == 'name':
          self.source_config[section]['name'] = value
        elif key == 'full_hostnames' and int(value) == 0:
          self.source_config[section]['full_hostnames'] = False
        elif key == 'max_mtime_days':
          self.source_config[section]['max_mtime_days'] = int(value)
        elif key == 'include_duplicates':
          self.source_config[section]['include_duplicates'] = bool(value)
        elif key == 'synthetic':
          self.source_config[section]['synthetic'] = bool(value)
        else:
          self.source_config[section]['search_paths'].add(value)

  def ListSourceTypes(self):
    """Get a list of all data sources we know about."""
    return sorted(self.source_config.keys())

  def ListSourcesWithDetails(self):
    """Get a list of all data sources found with total counts.

    Returns:
      List of tuples in form of (short_name, full_name, full_hosts, # of entries)
    """
    for source in self.ListSourceTypes():
      max_mtime = self.source_config[source]['max_mtime_days']
      self._GetHostsFromSource(source, min_file_size=MIN_FILE_SIZE,
                               max_mtime_age_days=max_mtime)

    details = []
    for source in self.source_cache:
      details.append((source,
                      self.source_config[source]['name'],
                      self.source_config[source]['synthetic'],
                      len(self.source_cache[source])))
    return sorted(details, key=lambda x: (x[2], x[3] * -1))

  def ListSourceTitles(self):
    """Return a list of sources in title + count format."""
    titles = []
    seen_synthetic = False
    seen_organic = False
    for (unused_type, name, is_synthetic, count) in self.ListSourcesWithDetails():
      if not is_synthetic:
        seen_organic = True

      if is_synthetic and seen_organic and not seen_synthetic:
        titles.append('-' * 36)
        seen_synthetic = True
      titles.append('%s (%s)' % (name, count))
    return titles

  def ConvertSourceTitleToType(self, detail):
    """Convert a detail name to a source type."""
    for source_type in self.source_config:
      if detail.startswith(self.source_config[source_type]['name']):
        return source_type

  def GetBestSourceDetails(self):
    return self.ListSourcesWithDetails()[0]

  def GetNameForSource(self, source):
    if source in self.source_config:
      return self.source_config[source]['name']
    else:
      # Most likely a custom file path
      return source

  def GetCachedRecordCountForSource(self, source):
    return len(self.source_cache[source])

  def _CreateRecordsFromHostEntries(self, entries, include_duplicates=False):
    """Create records from hosts, removing duplicate entries and IP's.

    Args:
      entries: A list of test-data entries.
      include_duplicates: Whether or not to filter duplicates (optional: False)

    Returns:
      A tuple of (filtered records, full_host_names (Boolean)

    Raises:
      ValueError: If no records could be grokked from the input.
    """
    last_entry = None

    records = []
    full_host_count = 0
    for entry in entries:
      if entry == last_entry and not include_duplicates:
        continue
      else:
        last_entry = entry

      if ' ' in entry:
        (record_type, host) = entry.split(' ')
      else:
        record_type = 'A'
        host = entry

      if not addr_util.IP_RE.match(host) and not addr_util.INTERNAL_RE.search(host):
        if not host.endswith('.'):
          host += '.'
        records.append((record_type, host))

        if addr_util.FQDN_RE.match(host):
          full_host_count += 1

    if not records:
      raise ValueError('No records could be created from: %s' % entries)

    # Now that we've read everything, are we dealing with domains or full hostnames?
    full_host_percent = full_host_count / float(len(records)) * 100
    if full_host_percent < MAX_FQDN_SYNTHESIZE_PERCENT:
      full_host_names = True
    else:
      full_host_names = False
    return (records, full_host_names)

  def GetTestsFromSource(self, source, count=50, select_mode=None):
    """Parse records from source, and return tuples to use for testing.

    Args:
      source: A source name (str) that has been configured.
      count: Number of tests to generate from the source (int)
      select_mode: automatic, weighted, random, chunk (str)

    Returns:
      A list of record tuples in the form of (req_type, hostname)

    Raises:
      ValueError: If no usable records are found from the data source.

    This is tricky because we support 3 types of input data:

    - List of domains
    - List of hosts
    - List of record_type + hosts
    """
    records = []

    if source in self.source_config:
      include_duplicates = self.source_config[source].get('include_duplicates', False)
    else:
      include_duplicates = False

    records = self._GetHostsFromSource(source)
    if not records:
      raise ValueError('Unable to generate records from %s (nothing found)' % source)

    self.msg('Generating tests from %s (%s records, selecting %s %s)'
             % (self.GetNameForSource(source), len(records), count, select_mode))
    (records, are_records_fqdn) = self._CreateRecordsFromHostEntries(records,
                                                                     include_duplicates=include_duplicates)
    # First try to resolve whether to use weighted or random.
    if select_mode in ('weighted', 'automatic', None):
      # If we are in include_duplicates mode (cachemiss, cachehit, etc.), we have different rules.
      if include_duplicates:
        if count > len(records):
          select_mode = 'random'
        else:
          select_mode = 'chunk'
      elif len(records) != len(set(records)):
        if select_mode == 'weighted':
          self.msg('%s data contains duplicates, switching select_mode to random' % source)
        select_mode = 'random'
      else:
        select_mode = 'weighted'

    self.msg('Selecting %s out of %s sanitized records (%s mode).' %
             (count, len(records), select_mode))
    if select_mode == 'weighted':
      records = selectors.WeightedDistribution(records, count)
    elif select_mode == 'chunk':
      records = selectors.ChunkSelect(records, count)
    elif select_mode == 'random':
      records = selectors.RandomSelect(records, count, include_duplicates=include_duplicates)
    else:
      raise ValueError('No such final selection mode: %s' % select_mode)

    # For custom filenames
    if source not in self.source_config:
      self.source_config[source] = {'synthetic': True}

    if are_records_fqdn:
      self.source_config[source]['full_hostnames'] = False
      self.msg('%s input appears to be predominantly domain names. Synthesizing FQDNs' % source)
      synthesized = []
      for (req_type, hostname) in records:
        if not addr_util.FQDN_RE.match(hostname):
          hostname = self._GenerateRandomHostname(hostname)
        synthesized.append((req_type, hostname))
      return synthesized
    else:
      return records

  def _GenerateRandomHostname(self, domain):
    """Generate a random hostname f or a given domain."""
    oracle = random.randint(0, 100)
    if oracle < 70:
      return 'www.%s' % domain
    elif oracle < 95:
      return domain
    elif oracle < 98:
      return 'static.%s' % domain
    else:
      return 'cache-%s.%s' % (random.randint(0, 10), domain)

  def _GetHostsFromSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Get data for a particular source. This needs to be fast.

    Args:
      source: A configured source type (str)
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      list of hostnames gathered from data source.

    The results of this function are cached by source type.
    """
    if source in self.source_cache:
      return self.source_cache[source]
    filename = self._FindBestFileForSource(source, min_file_size=min_file_size,
                                           max_mtime_age_days=max_mtime_age_days)
    if not filename:
      return None

    size_mb = os.path.getsize(filename) / 1024.0 / 1024.0
    self.msg('Reading %s: %s (%0.1fMB)' % (self.GetNameForSource(source), filename, size_mb))
    start_clock = DEFAULT_TIMER()
    if filename.endswith('.pcap') or filename.endswith('.tcp'):
      hosts = self._ExtractHostsFromPcapFile(filename)
    else:
      hosts = self._ExtractHostsFromHistoryFile(filename)

    if not hosts:
      hosts = self._ReadDataFile(filename)
    duration = DEFAULT_TIMER() - start_clock
    if duration > 5:
      self.msg('%s data took %1.1fs to read!' % (self.GetNameForSource(source), duration))
    self.source_cache[source] = hosts
    return hosts

  def _ExtractHostsFromHistoryFile(self, path):
    """Get a list of sanitized records from a history file containing URLs."""
    # This regexp is fairly general (no ip filtering), since we need speed more
    # than precision at this stage.
    parse_re = re.compile('https*://([\-\w]+\.[\-\w\.]+)')
    return parse_re.findall(open(path, 'rb').read())

  def _ExtractHostsFromPcapFile(self, path):
    """Get a list of requests out of a pcap file - requires tcpdump."""
    self.msg('Extracting requests from pcap file using tcpdump')
    cmd = 'tcpdump -r %s -n port 53' % path
    pipe = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).stdout
    parse_re = re.compile(' ([A-Z]+)\? ([\-\w\.]+)')
    requests = []
    for line in pipe:
      if '?' not in line:
        continue
      match = parse_re.search(line)
      if match:
        requests.append(' '.join(match.groups()))
    return requests

  def _ReadDataFile(self, path):
    """Read a line-based datafile."""
    records = []
    domains_re = re.compile('^\w[\w\.-]+[a-zA-Z]$')
    requests_re = re.compile('^[A-Z]{1,4} \w[\w\.-]+\.$')
    for line in open(path):
      if domains_re.match(line) or requests_re.match(line):
        records.append(line.rstrip())
    return records

  def _GetSourceSearchPaths(self, source):
    """Get a list of possible search paths (globs) for a given source."""

    # This is likely a custom file path
    if source not in self.source_config:
      return [source]

    search_paths = []
    environment_re = re.compile('%(\w+)%')

    # First get through resolving environment variables
    for path in self.source_config[source]['search_paths']:
      env_vars = set(environment_re.findall(path))
      if env_vars:
        for variable in env_vars:
          env_var = os.getenv(variable, False)
          if env_var:
            path = path.replace('%%%s%%' % variable, env_var)
          else:
            path = None

      # If everything is good, replace all '/'  chars with the os path variable.
      if path:
        path = path.replace('/', os.sep)
        search_paths.append(path)

        # This moment of weirdness brought to you by Windows XP(tm). If we find
        # a Local or Roaming keyword in path, add the other forms to the search
        # path.
        if sys.platform[:3] == 'win':
          keywords = ['Local', 'Roaming']
          for keyword in keywords:
            if keyword in path:
              replacement = keywords[keywords.index(keyword)-1]
              search_paths.append(path.replace('\\%s' % keyword, '\\%s' % replacement))
              search_paths.append(path.replace('\\%s' % keyword, ''))

    return search_paths

  def _FindBestFileForSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Find the best file (newest over X size) to use for a given source type.

    Args:
      source: source type
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      A file path.
    """
    found = []
    for path in self._GetSourceSearchPaths(source):
      if not os.path.isabs(path):
        path = util.FindDataFile(path)

      for filename in glob.glob(path):
        if min_file_size and os.path.getsize(filename) < min_file_size:
          self.msg('Skipping %s (only %sb)' % (filename, os.path.getsize(filename)))
        else:
          found.append(filename)

    if found:
      newest = sorted(found, key=os.path.getmtime)[-1]
      age_days = (time.time() - os.path.getmtime(newest)) / 86400
      if max_mtime_age_days and age_days > max_mtime_age_days:
        self.msg('Skipping %s (%2.0fd old)' % (newest, age_days))
      else:
        return newest
    else:
      return None

if __name__ == '__main__':
  parser = DataSources()
  print parser.ListSourceTypes()
  print parser.ListSourcesWithDetails()
  best = parser.ListSourcesWithDetails()[0][0]
  print len(parser.GetRecordsFromSource(best))
<|endoftext|>


<commit_before>from django.db import models
from django.utils import timezone
from django.contrib.auth.models import User

class Noticias(models.Model):
    STATUS_CHOICES = (
        ('draft', 'Rascunho'),
        ('published', 'Publicado'),
    )

    CATEGORY_CHOICES = (
        ('news', 'Noticia'),
    )

    title = models.CharField(max_length=140, verbose_name='Título')
    slug = models.SlugField(max_length=250, unique_for_date='publish')
    author = models.ForeignKey(User)
    url = models.URLField(max_length=300, verbose_name='Link/URL')
    summary = models.TextField(max_length=350, verbose_name='Resumo')
    category = models.CharField(max_length=30,
                              choices=CATEGORY_CHOICES,
                              default='news', verbose_name='Categoria')
    publish = models.DateTimeField(default=timezone.now,
                                   verbose_name='Publicação')
    created = models.DateTimeField(auto_now_add=True,
                                   verbose_name='Criação')
    updated = models.DateTimeField(auto_now=True,
                                   verbose_name='Atualização')
    status = models.CharField(max_length=10,
                              choices=STATUS_CHOICES,
                              default='draft', verbose_name='Status')

    def __str__(self):
        return self.title

    class Meta:
        verbose_name = 'Noticia'
        verbose_name_plural = 'Noticias'
        ordering = ('-publish',)
<commit_msg>utf-8<commit_after># coding: utf-8
from django.db import models
from django.utils import timezone
from django.contrib.auth.models import User

class Noticias(models.Model):
    STATUS_CHOICES = (
        ('draft', 'Rascunho'),
        ('published', 'Publicado'),
    )

    CATEGORY_CHOICES = (
        ('news', 'Noticia'),
    )

    title = models.CharField(max_length=140, verbose_name='Título')
    slug = models.SlugField(max_length=250, unique_for_date='publish')
    author = models.ForeignKey(User)
    url = models.URLField(max_length=300, verbose_name='Link/URL')
    summary = models.TextField(max_length=350, verbose_name='Resumo')
    category = models.CharField(max_length=30,
                              choices=CATEGORY_CHOICES,
                              default='news', verbose_name='Categoria')
    publish = models.DateTimeField(default=timezone.now,
                                   verbose_name='Publicação')
    created = models.DateTimeField(auto_now_add=True,
                                   verbose_name='Criação')
    updated = models.DateTimeField(auto_now=True,
                                   verbose_name='Atualização')
    status = models.CharField(max_length=10,
                              choices=STATUS_CHOICES,
                              default='draft', verbose_name='Status')

    def __str__(self):
        return self.title

    class Meta:
        verbose_name = 'Noticia'
        verbose_name_plural = 'Noticias'
        ordering = ('-publish',)
<|endoftext|>