<commit_before>import requests
import random
import time
import after_response
from datetime import datetime
from django.db.models import F, Count
from chef_buddy.models import Recipe, UserFlavorCompound, IngredientFlavorCompound
import os

_app_id = os.environ['APP_ID']
_app_key = os.environ['APP_KEY']

def speed_test(func):
    """wrapper function for testing speed"""
    def wrapper(*args, **kwargs):
        t1 = time.time()
        for x in range(1):
            results = func(*args, **kwargs)
        t2 = time.time()
        total_time = t2-t1
        print('\n')
        print('{} took {} seconds'.format(func.__name__, total_time))
        return results
    return wrapper

def get_recipes(amount):
    random_start = random.randint(1, (321961 - amount))
    all_recipes = requests.get('http://api.yummly.com/v1/api/recipes',
                                params={'_app_id':_app_id, '_app_key':_app_key,
                                        'maxResult':amount, 'requirePictures':'true',
                                        'start':random_start })
    return all_recipes.json()

@speed_test
def get_filtered_recipes(search, label_list, amount):
    params = get_parameters(search, label_list, amount)
    all_recipes = requests.get('http://api.yummly.com/v1/api/recipes', params=params)
    recipes = all_recipes.json()
    return recipes['matches']

@speed_test
def get_parameters(search, label_list, amount):
    """Takes in a list of dietary labels and then returns a parameter list to be sent to the yummly api"""
    query = search.replace(' ', '+')
    start = random.randint(1, 2000)
    if query:
        start = 1
    params = {'_app_id':_app_id, '_app_key':_app_key,
              'maxResult': amount, 'requirePictures':'true',
              'allowedDiet[]':[], 'allowedAllergy[]':[],
              'q':query, 'start':start}

    allergy_dict = {'dairy': '396^Dairy-Free', 'egg': '397^Egg-Free',
                    'gluten': '393^Gluten-Free', 'peanut': '394^Peanut-Free',
                    'seafood': '398^Seafood-Free', 'seasame': '399^Sesame-Free',
                    'soy': '400^Soy-Free', 'sulfite': '401^Sulfite-Free',
                    'nut': '395^Tree+Nut-Free', 'wheat': '392^Wheat-Free'}

    diet_dict = {'lacto': '388^Lacto vegetarian', 'ovo': '389^Ovo+vegetarian',
                 'pescetarian': '390^Pescetarian','vegan': '386^Vegan',
                 'vegetarian': '387^Lacto-ovo+vegetarian', 'paleo': '403^Paleo'}
    for label in label_list:
        if label in diet_dict:
            params['allowedDiet[]'].append(diet_dict[label])
        if label in allergy_dict:
            params['allowedAllergy[]'].append(allergy_dict[label])
    return params

def get_yummly_recipes():
    recipes = get_recipes(40)
    return recipes['matches']

def get_curated_random_recipes():
    category_list = ['chicken', 'beef', 'bread', 'tomato', 'cake', 'tuna', 'strawberry',
                     'mexican', 'italian', 'chinese', 'seafood', 'potato', 'cheese']
    params={'_app_id':_app_id, '_app_key':_app_key,
            'maxResult': 40, 'requirePictures':'true',
            'q':random.choice(category_list),
            'start':random.randint(1, 500)}
    all_recipes = requests.get('http://api.yummly.com/v1/api/recipes', params=params)
    recipes = all_recipes.json()
    return recipes['matches']

def recipe_ingr_parse(recipe_list):
    """recipe_list = raw list of recipes directly from yummly
    takes in a list of recipes and returns a dict of recipe_id, [ingredients]"""
    recipe_ingredient_dict = {}
    for recipe in recipe_list:
        recipe_ingredient_dict[recipe['id']]= [ingredient for ingredient in recipe['ingredients']]
    return recipe_ingredient_dict

def recipes_to_fc_id(recipe_list):
    """ recipe_list = Raw Recipe input from yummly
    Takes in a list of recipes and returns a dict of recipe_id to flavor compound id"""
    recipe_ingr_dict = recipe_ingr_parse(recipe_list)
    recipe_fc_dict = {}
    for recipe, ingredients in recipe_ingr_dict.items():
        recipe_fc_dict[recipe] = IngredientFlavorCompound.objects.filter(ingredient_id__in=ingredients)\
                                                                 .values_list('flavor_id', 'score')
    return recipe_fc_dict

@after_response.enable
def store_user_fc(user_id, recipe_id, taste):
    """Takes a user_id, recipe_id, and taste (should be a 1 or -1). It will lookup the recipe_id in the
    database to find associated food compounds. It will then lookup to see if the user has already liked
    disliked those flavor compounds. If they've liked them before, then it will update the score. If not,
    It will create a new row for that food compound."""
    user_id, taste = int(user_id), int(taste)
    if taste in [-1, 1]:
        flavor_compounds = Recipe.objects.filter(recipe_id=recipe_id).values_list('flavor_id', flat=True)
        exists = UserFlavorCompound.objects.filter(user_id=user_id, flavor_id__in=flavor_compounds)\
                                           .values_list('flavor_id', flat=True)
        UserFlavorCompound.objects.filter(user_id=user_id, flavor_id__in=exists). \
                                   update(score=F('score') + taste)
        new_fc = [num for num in set(flavor_compounds) if num not in set(exists)]
        UserFlavorCompound.objects.bulk_create([UserFlavorCompound(user_id=user_id,flavor_id=flavor,
                                                                   score=taste) for flavor in new_fc])
    return True

def recipe_id_to_object(recipe_id, recipe_list):
    """Looks recipe_id in recipe_list and returns recipe object"""
    for recipe in recipe_list:
        if recipe['id'] == recipe_id:
            return recipe

def user_to_recipe_counter(recipe_id_dict, user):
    """Takes in user's flavor compounds and recipe flavor compounds to produce a list of how
    many times the flavor compounds of the user appear in each recipe for the user. Each time
    a flavor compound appears, the score associated with the user's fc will be added to the recipe"""
    user = int(user)
    match_list = []
    for recipe_id, fc_score_list in recipe_id_dict.items():
        if len(fc_score_list) > 0:
            adj_scores, fc_list = in_common_fc_lookup(user, fc_score_list)
            all_user_fc = UserFlavorCompound.objects.filter(user_id=user).values_list('score', flat=True)
            score = calculate_recipe_score(fc_list, adj_scores, all_user_fc)
        else:
            score = 0
        match_list.append((recipe_id, score))
    return match_list

def in_common_fc_lookup(user, fc_score_list):
    fc_list = [x[0] for x in fc_score_list]
    in_common = dict(UserFlavorCompound.objects.values_list('flavor_id', 'score'). \
                                                filter(user_id=user, flavor_id__in=list(fc_list)))
    adjusted_score_list = [(in_common[fc_id]*score) for fc_id, score in fc_score_list if fc_id in in_common.keys()]
    return adjusted_score_list, fc_list

def calculate_recipe_score(recipe_fc_list, user_fc_scores, all_user_fc):
    total_fc_score = sum([abs(fc) for fc in all_user_fc])
    normalized_scoring = [(fc/total_fc_score) for fc in user_fc_scores]
    engine_score = (sum(normalized_scoring) / len(recipe_fc_list))
    return engine_score

def user_shown_score(recipe_fc_list, user):
    recipe_fc_list = [x[0] for x in recipe_fc_list]
    in_common_fc_score = UserFlavorCompound.objects. \
                         filter(user_id=user, score__gt=1, flavor_id__in=list(recipe_fc_list)). \
                         values_list('flavor_id')
    all_user_fc = UserFlavorCompound.objects.filter(user_id=user).values_list('flavor_id')
    if (len(all_user_fc) < 40):
        return '?'
    else:
        total_compounds = len(recipe_fc_list)
        if total_compounds == 0:
            return 1
        user_score = (len(in_common_fc_score) / len(recipe_fc_list)) * 100
        return round_float(user_score)

def large_image(json):
    """
    Takes the image from the Yummly api and returns a larger image by changing the URL.
    :return: json object
    """
    image = json["imageUrlsBySize"]['90']
    json['largeImage'] = image.replace('=s90-c', '=s600')
    return json

def store_recipe_fc(recipe_id, flavor_compounds):
    """recipe_id = id of the recipe needing to be housed
    flavor_compounds = list of flavor compounds associated with recipe
    This function is designed to take the above variables and store them in separate rows in the db"""
    if not Recipe.objects.filter(recipe_id=recipe_id):
        recipe_list = [Recipe(recipe_id=str(recipe_id), flavor_id=int(fc_id)) for fc_id, score in flavor_compounds]
        Recipe.objects.bulk_create(recipe_list)
    return True

@after_response.enable
def log_recommendation(dict_of_logs):
    with open('chef_buddy/raw_data/rec_log.txt', 'a') as the_file:
        the_file.write(str(datetime.now()))
        the_file.write('\n')
        for header, log in dict_of_logs.items():
            the_file.write(header)
            the_file.write('\n\n')
            the_file.write(str(log))
            the_file.write('\n\n')
        the_file.write('\n\n\n\n\n')
    return True

def round_float(float):
    """iround(number) -> integer
    Round a number to the nearest integer."""
    y = round(float) - .5
    return int(y) + (y > 0)
<commit_msg>changed large image to resize'<commit_after>import requests
import random
import time
import after_response
from datetime import datetime
from django.db.models import F, Count
from chef_buddy.models import Recipe, UserFlavorCompound, IngredientFlavorCompound
import os

_app_id = os.environ['APP_ID']
_app_key = os.environ['APP_KEY']

def speed_test(func):
    """wrapper function for testing speed"""
    def wrapper(*args, **kwargs):
        t1 = time.time()
        for x in range(1):
            results = func(*args, **kwargs)
        t2 = time.time()
        total_time = t2-t1
        print('\n')
        print('{} took {} seconds'.format(func.__name__, total_time))
        return results
    return wrapper

def get_recipes(amount):
    random_start = random.randint(1, (321961 - amount))
    all_recipes = requests.get('http://api.yummly.com/v1/api/recipes',
                                params={'_app_id':_app_id, '_app_key':_app_key,
                                        'maxResult':amount, 'requirePictures':'true',
                                        'start':random_start })
    return all_recipes.json()

@speed_test
def get_filtered_recipes(search, label_list, amount):
    params = get_parameters(search, label_list, amount)
    all_recipes = requests.get('http://api.yummly.com/v1/api/recipes', params=params)
    recipes = all_recipes.json()
    return recipes['matches']

@speed_test
def get_parameters(search, label_list, amount):
    """Takes in a list of dietary labels and then returns a parameter list to be sent to the yummly api"""
    query = search.replace(' ', '+')
    start = random.randint(1, 2000)
    if query:
        start = 1
    params = {'_app_id':_app_id, '_app_key':_app_key,
              'maxResult': amount, 'requirePictures':'true',
              'allowedDiet[]':[], 'allowedAllergy[]':[],
              'q':query, 'start':start}

    allergy_dict = {'dairy': '396^Dairy-Free', 'egg': '397^Egg-Free',
                    'gluten': '393^Gluten-Free', 'peanut': '394^Peanut-Free',
                    'seafood': '398^Seafood-Free', 'seasame': '399^Sesame-Free',
                    'soy': '400^Soy-Free', 'sulfite': '401^Sulfite-Free',
                    'nut': '395^Tree+Nut-Free', 'wheat': '392^Wheat-Free'}

    diet_dict = {'lacto': '388^Lacto vegetarian', 'ovo': '389^Ovo+vegetarian',
                 'pescetarian': '390^Pescetarian','vegan': '386^Vegan',
                 'vegetarian': '387^Lacto-ovo+vegetarian', 'paleo': '403^Paleo'}
    for label in label_list:
        if label in diet_dict:
            params['allowedDiet[]'].append(diet_dict[label])
        if label in allergy_dict:
            params['allowedAllergy[]'].append(allergy_dict[label])
    return params

def get_yummly_recipes():
    recipes = get_recipes(40)
    return recipes['matches']

def get_curated_random_recipes():
    category_list = ['chicken', 'beef', 'bread', 'tomato', 'cake', 'tuna', 'strawberry',
                     'mexican', 'italian', 'chinese', 'seafood', 'potato', 'cheese']
    params={'_app_id':_app_id, '_app_key':_app_key,
            'maxResult': 40, 'requirePictures':'true',
            'q':random.choice(category_list),
            'start':random.randint(1, 500)}
    all_recipes = requests.get('http://api.yummly.com/v1/api/recipes', params=params)
    recipes = all_recipes.json()
    return recipes['matches']

def recipe_ingr_parse(recipe_list):
    """recipe_list = raw list of recipes directly from yummly
    takes in a list of recipes and returns a dict of recipe_id, [ingredients]"""
    recipe_ingredient_dict = {}
    for recipe in recipe_list:
        recipe_ingredient_dict[recipe['id']]= [ingredient for ingredient in recipe['ingredients']]
    return recipe_ingredient_dict

def recipes_to_fc_id(recipe_list):
    """ recipe_list = Raw Recipe input from yummly
    Takes in a list of recipes and returns a dict of recipe_id to flavor compound id"""
    recipe_ingr_dict = recipe_ingr_parse(recipe_list)
    recipe_fc_dict = {}
    for recipe, ingredients in recipe_ingr_dict.items():
        recipe_fc_dict[recipe] = IngredientFlavorCompound.objects.filter(ingredient_id__in=ingredients)\
                                                                 .values_list('flavor_id', 'score')
    return recipe_fc_dict

@after_response.enable
def store_user_fc(user_id, recipe_id, taste):
    """Takes a user_id, recipe_id, and taste (should be a 1 or -1). It will lookup the recipe_id in the
    database to find associated food compounds. It will then lookup to see if the user has already liked
    disliked those flavor compounds. If they've liked them before, then it will update the score. If not,
    It will create a new row for that food compound."""
    user_id, taste = int(user_id), int(taste)
    if taste in [-1, 1]:
        flavor_compounds = Recipe.objects.filter(recipe_id=recipe_id).values_list('flavor_id', flat=True)
        exists = UserFlavorCompound.objects.filter(user_id=user_id, flavor_id__in=flavor_compounds)\
                                           .values_list('flavor_id', flat=True)
        UserFlavorCompound.objects.filter(user_id=user_id, flavor_id__in=exists). \
                                   update(score=F('score') + taste)
        new_fc = [num for num in set(flavor_compounds) if num not in set(exists)]
        UserFlavorCompound.objects.bulk_create([UserFlavorCompound(user_id=user_id,flavor_id=flavor,
                                                                   score=taste) for flavor in new_fc])
    return True

def recipe_id_to_object(recipe_id, recipe_list):
    """Looks recipe_id in recipe_list and returns recipe object"""
    for recipe in recipe_list:
        if recipe['id'] == recipe_id:
            return recipe

def user_to_recipe_counter(recipe_id_dict, user):
    """Takes in user's flavor compounds and recipe flavor compounds to produce a list of how
    many times the flavor compounds of the user appear in each recipe for the user. Each time
    a flavor compound appears, the score associated with the user's fc will be added to the recipe"""
    user = int(user)
    match_list = []
    for recipe_id, fc_score_list in recipe_id_dict.items():
        if len(fc_score_list) > 0:
            adj_scores, fc_list = in_common_fc_lookup(user, fc_score_list)
            all_user_fc = UserFlavorCompound.objects.filter(user_id=user).values_list('score', flat=True)
            score = calculate_recipe_score(fc_list, adj_scores, all_user_fc)
        else:
            score = 0
        match_list.append((recipe_id, score))
    return match_list

def in_common_fc_lookup(user, fc_score_list):
    fc_list = [x[0] for x in fc_score_list]
    in_common = dict(UserFlavorCompound.objects.values_list('flavor_id', 'score'). \
                                                filter(user_id=user, flavor_id__in=list(fc_list)))
    adjusted_score_list = [(in_common[fc_id]*score) for fc_id, score in fc_score_list if fc_id in in_common.keys()]
    return adjusted_score_list, fc_list

def calculate_recipe_score(recipe_fc_list, user_fc_scores, all_user_fc):
    total_fc_score = sum([abs(fc) for fc in all_user_fc])
    normalized_scoring = [(fc/total_fc_score) for fc in user_fc_scores]
    engine_score = (sum(normalized_scoring) / len(recipe_fc_list))
    return engine_score

def user_shown_score(recipe_fc_list, user):
    recipe_fc_list = [x[0] for x in recipe_fc_list]
    in_common_fc_score = UserFlavorCompound.objects. \
                         filter(user_id=user, score__gt=1, flavor_id__in=list(recipe_fc_list)). \
                         values_list('flavor_id')
    all_user_fc = UserFlavorCompound.objects.filter(user_id=user).values_list('flavor_id')
    if (len(all_user_fc) < 40):
        return '?'
    else:
        total_compounds = len(recipe_fc_list)
        if total_compounds == 0:
            return 1
        user_score = (len(in_common_fc_score) / len(recipe_fc_list)) * 100
        return round_float(user_score)

def large_image(json):
    """
    Takes the image from the Yummly api and returns a larger image by changing the URL.
    :return: json object
    """
    image = json["imageUrlsBySize"]['90']
    json['largeImage'] = image.replace('=s90-c', '=s600-c')
    return json

def store_recipe_fc(recipe_id, flavor_compounds):
    """recipe_id = id of the recipe needing to be housed
    flavor_compounds = list of flavor compounds associated with recipe
    This function is designed to take the above variables and store them in separate rows in the db"""
    if not Recipe.objects.filter(recipe_id=recipe_id):
        recipe_list = [Recipe(recipe_id=str(recipe_id), flavor_id=int(fc_id)) for fc_id, score in flavor_compounds]
        Recipe.objects.bulk_create(recipe_list)
    return True

@after_response.enable
def log_recommendation(dict_of_logs):
    with open('chef_buddy/raw_data/rec_log.txt', 'a') as the_file:
        the_file.write(str(datetime.now()))
        the_file.write('\n')
        for header, log in dict_of_logs.items():
            the_file.write(header)
            the_file.write('\n\n')
            the_file.write(str(log))
            the_file.write('\n\n')
        the_file.write('\n\n\n\n\n')
    return True

def round_float(float):
    """iround(number) -> integer
    Round a number to the nearest integer."""
    y = round(float) - .5
    return int(y) + (y > 0)
<|endoftext|>


<commit_before># Copyright (c) 2015, The MITRE Corporation. All rights reserved.
# See LICENSE.txt for complete terms.

from StringIO import StringIO
import unittest

from stix.incident import Incident
import stix.bindings.incident as incident_binding
from cybox.common import StructuredText
from stix.test import EntityTestCase

INCIDENT_CATEGORIES = """<?xml version="1.0" encoding="UTF-8"?>
<incident:Incident
 xmlns:incident="http://stix.mitre.org/Incident-1"
 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
 xsi:schemaLocation="http://stix.mitre.org/Incident-1 http://stix.mitre.org/XMLSchema/incident/1.0.1/incident.xsd">
 <incident:Categories>
  <incident:Category>Foo</incident:Category>
  <incident:Category>Bar</incident:Category>
 </incident:Categories>
</incident:Incident>
"""

class IncidentTest(EntityTestCase, unittest.TestCase):
    klass = Incident
    _full_dict = {
'attributed_threat_actors': {'scope': 'exclusive',
                             'threat_actors': [{'threat_actor': {'description': 'A Threat Actor Description',
                                                                  'id': 'example:threatactor-1',
                                                                  'sophistications': [{'value': {"value" : "Novice",
                                                                                "xsi:type" : "stixVocabs:ThreatActorSophisticationVocab-1.0"}}],
                                                                  'title': 'A Threat Actor',
                                                                  'version': '1.1'}}]},
'categories': [{'value': 'Unauthorized Access',
                 'xsi:type': 'stixVocabs:IncidentCategoryVocab-1.0'},
                {'value': 'Improper Usage',
                 'xsi:type': 'stixVocabs:IncidentCategoryVocab-1.0'}],
'coa_taken': [{'course_of_action': {'timestamp': '2014-05-05T14:50:25.992383+00:00', 'version': '1.1', 'id': 'example:coa-74e50620-7536-4fcd-94db-a6889f75e098'}},
              {'course_of_action': {'timestamp': '2014-05-05T14:50:25.992384+00:00', 'version': '1.1', 'id': 'example:coa-74e50620-7536-4fcd-94db-a6889f75e099'}}],
'coordinators': [{'description': "Mr. Evil's enemy",
                   'identity': {'name': 'Ms. Coordinator'}}],
'description': 'The Datacenter was broken into.',
'external_ids': [{'source': 'some source',
                   'value': '478392-feb3ca-98a9ef-984392742'}],
'impact_assessment': {'direct_impact_summary': {'asset_losses': {'value': 'Minor',
                                                                  'xsi:type': 'stixVocabs:ImpactRatingVocab-1.0'},
                                                 'business_mission_disruption': {'value': 'Major',
                                                                                'xsi:type': 'stixVocabs:ImpactRatingVocab-1.0'},
                                                 'response_and_recovery_costs': {'value': 'Moderate',
                                                                                'xsi:type': 'stixVocabs:ImpactRatingVocab-1.0'}},
                       'effects': [{'value': 'User Data Loss',
                                                'xsi:type': 'stixVocabs:IncidentEffectVocab-1.0'},
                                               {'value': 'Data Breach or Compromise',
                                                'xsi:type': 'stixVocabs:IncidentEffectVocab-1.0'}],
                       'impact_qualification': {'value': 'Catastrophic',
                                                'xsi:type': 'stixVocabs:ImpactQualificationVocab-1.0'},
                       'indirect_impact_summary': {'brand_and_market_damage': {'value': 'No',
                                                                               'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'},
                                                   'increased_operating_costs': {'value': 'No',
                                                                                'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'},
                                                   'legal_and_regulatory_costs': {'value': 'Unknown',
                                                                                'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'},
                                                   'loss_of_competitive_advantage': {'value': 'Yes',
                                                                                'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'}},
                       'total_loss_estimation': {'actual_total_loss_estimation': {'amount': '50.45',
                                                                                'iso_currency_code': 'USD'},
                                                 'initial_reported_total_loss_estimation': {'amount': '99.99',
                                                                                'iso_currency_code': 'USD'}}},
 'leveraged_ttps': {'scope': 'inclusive',
                    'ttps': [{'confidence': {'value': {'value': 'Medium',
                                                       'xsi:type': 'stixVocabs:HighMediumLowVocab-1.0'}},
                              'ttp': {'id': 'example:TTP-1',
                                      'version': '1.1'}}]},
 'related_indicators': {'indicators': [{'indicator': {'description': 'An indicator containing a File observable with an associated hash',
                                                      'id': 'example:indicator-1ae45e9c-9b0b-11e3-ada0-28cfe912ced6',
                                                      'observable': {'id': 'example:Observable-fdaa7cec-f8be-494d-b83f-575f6f018666',
                                                                     'object': {'id': 'example:File-ec52e6bc-2d7e-44e2-911b-468bb775a5c6',
                                                                                'properties': {'hashes': [{'simple_hash_value': u'4EC0027BEF4D7E1786A04D021FA8A67F',
                                                                                'type': u'MD5'}],
                                                                                'xsi:type': 'FileObjectType'}}},
                                                      'producer': {'description': 'A sample description',
                                                                   'identity': {'id': 'example:Identity-1ae603ab-9b0b-11e3-980e-28cfe912ced8',
                                                                                'specification': {'party_name': {'name_lines': [{'value': 'Foo'},
                                                                                {'value': 'Bar'}],
                                                                                'person_names': [{'name_elements': [{'value': 'John Smith'}]},
                                                                                {'name_elements': [{'value': 'Jill Smith'}]}]}},
                                                                                'xsi:type': 'ciqIdentity:CIQIdentity3.0InstanceType'},
                                                                   'time': {'produced_time': '2014-02-21T10:16:14.947201'}},
                                                      'title': 'File Hash Example',
                                                      'version': '2.1.1'}}],
                        'scope': 'exclusive'},
 'reporter': {'description': "Mr. Evil's enemy",
              'identity': {'name': 'Ms. Good'}},
 'responders': [{'description': "Mr. Evil's enemy",
                 'identity': {'name': 'Ms. Responder'}}],
 'time': {'containment_achieved': '2005-02-21T10:25:10.894398',
          'first_data_exfiltration': '2002-02-21T10:25:10.894398',
          'first_malicious_action': '2000-02-21T10:25:10.894398',
          'incident_closed': '2008-02-21T10:25:10.894398',
          'incident_discovery': '2003-02-21T10:25:10.894398',
          'incident_opened': '2004-02-21T10:25:10.894398',
          'incident_reported': '2007-02-21T10:25:10.894398',
          'initial_compromise': '2001-02-21T10:25:10.894398',
          'restoration_achieved': '2006-02-21T10:25:10.894398'},
 'title': 'Datacenter Compromise',
 'version': '1.1',
 'victims': [{'name': 'John Smith'},
             {'id': 'example:Identity-1ae603ab-9b0b-11e3-980e-28cfe912ced6'}],
 'information_source': {
              'description': "Mr. Evil's enemy",
              'identity': {
                  'name': "Ms. Good",
              },
          },
'security_compromise': {
    "value": "Suspected",
    "xsi:type":"stixVocabs:SecurityCompromiseVocab-1.0"
},
'history': {
    'history_items':
        [
            {
                'journal_entry': {
                    'value': 'hi',
                    'author': 'Paul'
                }
            }
        ]
},
'affected_assets': [
    {
        'business_function_or_role': 'Foobar',
        'description': 'Foobar',
        'location_class': {
            'value': 'Unknown',
            'xsi:type': 'stixVocabs:LocationClassVocab-1.0'
        },
        'management_class': {
            'value': 'Unknown',
            'xsi:type': 'stixVocabs:ManagementClassVocab-1.0'
        },
        'nature_of_security_effect': [
            {
                'description_of_effect': 'Foobar',
                'duration_of_availability_loss': {
                    'value': 'Days',
                    'xsi:type': 'stixVocabs:LossDurationVocab-1.0'
                },
                'non_public_data_compromised': {
                    'value': 'Yes',
                    'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0',
                    'data_encrypted': True
                },
                'type_of_availability_loss': {
                    'value': 'Loss',
                    'xsi:type': 'stixVocabs:AvailabilityLossTypeVocab-1.1.1'
                }
            }
        ],
        'ownership_class': {
            'value': 'Unknown',
            'xsi:type': 'stixVocabs:OwnershipClassVocab-1.0'
        },
        'type': {'value': 'Foobar'}
    }
]
}

    def test_parse_category(self):
        incident = incident_binding.parseString(INCIDENT_CATEGORIES)
        self.assertTrue(incident is not None)
        self.assertEqual(2, len(incident.Categories.Category))

        categories = incident.Categories.Category
        self.assertEqual('Foo', categories[0].valueOf_)
        self.assertEqual('Bar', categories[1].valueOf_)

    def test_description_output(self):
        incident = incident_binding.IncidentType()

        assets = incident_binding.AffectedAssetsType()
        asset = incident_binding.AffectedAssetType()

        description = StructuredText("A Description")
        asset.Structured_Description = description.to_obj()

        assets.add_Affected_Asset(asset)
        incident.Affected_Assets = assets

        s = StringIO()

        incident.export(s.write, 0, {'http://stix.mitre.org/Incident-1': 'incident'})
        xml = s.getvalue()
        self.assertTrue("A Description" in xml, "Description not exported")

if __name__ == "__main__":
    unittest.main()
<commit_msg>WIP: Adding/refactoring Incident unit tests [ci skip]<commit_after># Copyright (c) 2015, The MITRE Corporation. All rights reserved.
# See LICENSE.txt for complete terms.

import unittest
import StringIO

from cybox.common import StructuredText

from stix.test import EntityTestCase

import stix.common.vocabs as vocabs
import stix.incident as incident
import stix.bindings.incident as incident_binding

INCIDENT_CATEGORIES = """<?xml version="1.0" encoding="UTF-8"?>
<incident:Incident
 xmlns:incident="http://stix.mitre.org/Incident-1"
 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
 xsi:schemaLocation="http://stix.mitre.org/Incident-1 http://stix.mitre.org/XMLSchema/incident/1.0.1/incident.xsd">
 <incident:Categories>
  <incident:Category>Foo</incident:Category>
  <incident:Category>Bar</incident:Category>
 </incident:Categories>
</incident:Incident>
"""

class RT(object):
    def test_round_trip_rt(self):
        if type(self) == type(RT):
            return

        obj = self.klass.from_dict(self._full_dict)
        dict2 = obj.to_dict()
        self.assertEqual(self._full_dict, dict2)

class IncidentTest(EntityTestCase, unittest.TestCase):
    klass = incident.Incident
    _full_dict = {
        'id': 'example:test-1'
    }

    def test_base(self):
        d = {
            'version': '1.1.1',
            'title': 'Test Title',
            'description': 'Test Description',
            'short_description': "Test Short Description",
            'timestamp': '2015-03-06T14:35:23.375304+00:00',
        }
        self._test_partial_dict(d)

    def test_external_ids(self):
        d = {
            'external_ids': [
                {
                    'source': 'foo',
                    'value': '478392-feb3ca-98a9ef-984392742'
                },
                {
                    'source': 'bar',
                    'value': '478392-feb3ca-98a9ef-984392742'
                },
            ],
        }
        self._test_partial_dict(d)

    def test_time(self):
        d = {
            'time': TimeTest._full_dict
        }
        self._test_partial_dict(d)

    def test_categories(self):
        d = {
            'categories': CategoriesTest._full_dict
        }
        self._test_partial_dict(d)

    def test_reporter(self):
        # This is an InformationSource instance, which is tested in the
        # stix.test.common package.
        d = {
            'reporter': {
                'description': 'Test',
                'identity': {
                    'name': 'Spooderman'
                }
            }
        }
        self._test_partial_dict(d)

    def test_responder(self):
        d = {'responders': InformationSourcesTest._full_dict}
        self._test_partial_dict(d)

    def test_coordinators(self):
        d = {'coordinators': InformationSourcesTest._full_dict}
        self._test_partial_dict(d)

    def test_victims(self):
        d = {'victims': VictimsTest._full_dict}
        self._test_partial_dict(d)

class VictimsTest(unittest.TestCase, RT):
    klass = incident._Victims

    _full_dict = [
        {'name': 'Spooderman'},
        {'name': 'Spooderman'}
    ]

class TimeTest(EntityTestCase, unittest.TestCase):
    klass = incident.Time
    _full_dict = {
            'containment_achieved': '2005-02-21T10:25:10.894398',
            'first_data_exfiltration': '2002-02-21T10:25:10.894398',
            'first_malicious_action': '2000-02-21T10:25:10.894398',
            'incident_closed': '2008-02-21T10:25:10.894398',
            'incident_discovery': '2003-02-21T10:25:10.894398',
            'incident_opened': '2004-02-21T10:25:10.894398',
            'incident_reported': '2007-02-21T10:25:10.894398',
            'initial_compromise': '2001-02-21T10:25:10.894398',
            'restoration_achieved': '2006-02-21T10:25:10.894398'
    }

class CategoriesTest(EntityTestCase, unittest.TestCase):
    klass = incident.IncidentCategories

    _full_dict = [
        {
            'value': vocabs.IncidentCategory.TERM_DENIAL_OF_SERVICE,
            'xsi:type': vocabs.IncidentCategory._XSI_TYPE
        },
        {
            'value': vocabs.IncidentCategory.TERM_IMPROPER_USAGE,
            'xsi:type': vocabs.IncidentCategory._XSI_TYPE
        }
    ]

class InformationSourcesTest(unittest.TestCase, RT):
    klass = incident._InformationSources

    _full_dict = [
        {
            'description': 'Test',
            'identity': {
                'name': 'Spooderman'
            }
        },
        {
            'description': 'Test',
            'identity': {
                'name': 'Spooderman'
            }
        }
    ]

    d = {
        'attributed_threat_actors': {'scope': 'exclusive',
                             'threat_actors': [{'threat_actor': {'description': 'A Threat Actor Description',
                                                                  'id': 'example:threatactor-1',
                                                                  'sophistications': [{'value': {"value" : "Novice",
                                                                                "xsi:type" : "stixVocabs:ThreatActorSophisticationVocab-1.0"}}],
                                                                  'title': 'A Threat Actor',
                                                                  'version': '1.1'}}]},
        'coa_taken': [{'course_of_action': {'timestamp': '2014-05-05T14:50:25.992383+00:00', 'version': '1.1', 'id': 'example:coa-74e50620-7536-4fcd-94db-a6889f75e098'}},
                      {'course_of_action': {'timestamp': '2014-05-05T14:50:25.992384+00:00', 'version': '1.1', 'id': 'example:coa-74e50620-7536-4fcd-94db-a6889f75e099'}}],
        'coordinators': [{'description': "Mr. Evil's enemy",
                           'identity': {'name': 'Ms. Coordinator'}}],

        'impact_assessment': {'direct_impact_summary': {'asset_losses': {'value': 'Minor',
                                                                          'xsi:type': 'stixVocabs:ImpactRatingVocab-1.0'},
                                                         'business_mission_disruption': {'value': 'Major',
                                                                                'xsi:type': 'stixVocabs:ImpactRatingVocab-1.0'},
                                                         'response_and_recovery_costs': {'value': 'Moderate',
                                                                                'xsi:type': 'stixVocabs:ImpactRatingVocab-1.0'}},
                               'effects': [{'value': 'User Data Loss',
                                                        'xsi:type': 'stixVocabs:IncidentEffectVocab-1.0'},
                                                       {'value': 'Data Breach or Compromise',
                                                        'xsi:type': 'stixVocabs:IncidentEffectVocab-1.0'}],
                               'impact_qualification': {'value': 'Catastrophic',
                                                        'xsi:type': 'stixVocabs:ImpactQualificationVocab-1.0'},
                               'indirect_impact_summary': {'brand_and_market_damage': {'value': 'No',
                                                                                'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'},
                                                           'increased_operating_costs': {'value': 'No',
                                                                                'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'},
                                                           'legal_and_regulatory_costs': {'value': 'Unknown',
                                                                                'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'},
                                                           'loss_of_competitive_advantage': {'value': 'Yes',
                                                                                'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0'}},
                               'total_loss_estimation': {'actual_total_loss_estimation': {'amount': '50.45',
                                                                                'iso_currency_code': 'USD'},
                                                         'initial_reported_total_loss_estimation': {'amount': '99.99',
                                                                                'iso_currency_code': 'USD'}}},
         'leveraged_ttps': {'scope': 'inclusive',
                            'ttps': [{'confidence': {'value': {'value': 'Medium',
                                                               'xsi:type': 'stixVocabs:HighMediumLowVocab-1.0'}},
                                      'ttp': {'id': 'example:TTP-1',
                                              'version': '1.1'}}]},
         'related_indicators': {'indicators': [{'indicator': {'description': 'An indicator containing a File observable with an associated hash',
                                                              'id': 'example:indicator-1ae45e9c-9b0b-11e3-ada0-28cfe912ced6',
                                                              'observable': {'id': 'example:Observable-fdaa7cec-f8be-494d-b83f-575f6f018666',
                                                                             'object': {'id': 'example:File-ec52e6bc-2d7e-44e2-911b-468bb775a5c6',
                                                                                'properties': {'hashes': [{'simple_hash_value': u'4EC0027BEF4D7E1786A04D021FA8A67F',
                                                                                'type': u'MD5'}],
                                                                                'xsi:type': 'FileObjectType'}}},
                                                              'producer': {'description': 'A sample description',
                                                                           'identity': {'id': 'example:Identity-1ae603ab-9b0b-11e3-980e-28cfe912ced8',
                                                                                'specification': {'party_name': {'name_lines': [{'value': 'Foo'},
                                                                                {'value': 'Bar'}],
                                                                                'person_names': [{'name_elements': [{'value': 'John Smith'}]},
                                                                                {'name_elements': [{'value': 'Jill Smith'}]}]}},
                                                                                'xsi:type': 'ciqIdentity:CIQIdentity3.0InstanceType'},
                                                                           'time': {'produced_time': '2014-02-21T10:16:14.947201'}},
                                                              'title': 'File Hash Example',
                                                              'version': '2.1.1'}}],
                                'scope': 'exclusive'},
         'reporter': {'description': "Mr. Evil's enemy",
                      'identity': {'name': 'Ms. Good'}},
         'responders': [{'description': "Mr. Evil's enemy",
                         'identity': {'name': 'Ms. Responder'}}],
         'victims': [{'name': 'John Smith'},
                     {'id': 'example:Identity-1ae603ab-9b0b-11e3-980e-28cfe912ced6'}],
         'information_source': {
                      'description': "Mr. Evil's enemy",
                      'identity': {
                          'name': "Ms. Good",
                      },
                  },
        'security_compromise': {
            "value": "Suspected",
            "xsi:type":"stixVocabs:SecurityCompromiseVocab-1.0"
        },
        'history': {
            'history_items':
                [
                    {
                        'journal_entry': {
                            'value': 'hi',
                            'author': 'Paul'
                        }
                    }
                ]
        },
        'affected_assets': [
            {
                'business_function_or_role': 'Foobar',
                'description': 'Foobar',
                'location_class': {
                    'value': 'Unknown',
                    'xsi:type': 'stixVocabs:LocationClassVocab-1.0'
                },
                'management_class': {
                    'value': 'Unknown',
                    'xsi:type': 'stixVocabs:ManagementClassVocab-1.0'
                },
                'nature_of_security_effect': [
                    {
                        'description_of_effect': 'Foobar',
                        'duration_of_availability_loss': {
                            'value': 'Days',
                            'xsi:type': 'stixVocabs:LossDurationVocab-1.0'
                        },
                        'non_public_data_compromised': {
                            'value': 'Yes',
                            'xsi:type': 'stixVocabs:SecurityCompromiseVocab-1.0',
                            'data_encrypted': True
                        },
                        'type_of_availability_loss': {
                            'value': 'Loss',
                            'xsi:type': 'stixVocabs:AvailabilityLossTypeVocab-1.1.1'
                        }
                    }
                ],
                'ownership_class': {
                    'value': 'Unknown',
                    'xsi:type': 'stixVocabs:OwnershipClassVocab-1.0'
                },
                'type': {'value': 'Foobar'}
            }
        ]
}

    def test_parse_category(self):
        incident = incident_binding.parseString(INCIDENT_CATEGORIES)
        self.assertTrue(incident is not None)
        self.assertEqual(2, len(incident.Categories.Category))

        categories = incident.Categories.Category
        self.assertEqual('Foo', categories[0].valueOf_)
        self.assertEqual('Bar', categories[1].valueOf_)

    def test_description_output(self):
        incident = incident_binding.IncidentType()

        assets = incident_binding.AffectedAssetsType()
        asset = incident_binding.AffectedAssetType()

        description = StructuredText("A Description")
        asset.Structured_Description = description.to_obj()

        assets.add_Affected_Asset(asset)
        incident.Affected_Assets = assets

        s = StringIO.StringIO()

        incident.export(s.write, 0, {'http://stix.mitre.org/Incident-1': 'incident'})
        xml = s.getvalue()
        self.assertTrue("A Description" in xml, "Description not exported")

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


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

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

import sys
import time
import socket
import errno
import logging

info = sys.version_info
if not (info[0] == 2 and info[1] >= 7):
    print 'Python 2.7 required'
    sys.exit(1)

import argparse
from shadowsocks import eventloop, asyncdns, lru_cache

BUF_SIZE = 16384

CACHE_TIMEOUT = 10

GFW_LIST = set(["74.125.127.102", "74.125.155.102", "74.125.39.102",
                "74.125.39.113", "209.85.229.138", "128.121.126.139",
                "159.106.121.75", "169.132.13.103", "192.67.198.6",
                "202.106.1.2", "202.181.7.85", "203.161.230.171",
                "203.98.7.65", "207.12.88.98", "208.56.31.43",
                "209.145.54.50", "209.220.30.174", "209.36.73.33",
                "211.94.66.147", "213.169.251.35", "216.221.188.182",
                "216.234.179.13", "243.185.187.39", "37.61.54.158",
                "4.36.66.178", "46.82.174.68", "59.24.3.173", "64.33.88.161",
                "64.33.99.47", "64.66.163.251", "65.104.202.252",
                "65.160.219.113", "66.45.252.237", "72.14.205.104",
                "72.14.205.99", "78.16.49.15", "8.7.198.45", "93.46.8.89"])

class DNSRelay(object):
    def __init__(self, config):
        self._loop = None
        self._config = config
        self._last_time = time.time()

        self._local_addr = (config['local_address'], 53)
        self._remote_addr = (config['dns'], 53)

    def add_to_loop(self, loop):
        if self._loop:
            raise Exception('already add to loop')
        self._loop = loop
        loop.add_handler(self.handle_events)

    def handle_events(self, events):
        pass

class UDPDNSRelay(DNSRelay):
    def __init__(self, config):
        DNSRelay.__init__(self, config)

        self._id_to_addr = lru_cache.LRUCache(CACHE_TIMEOUT)
        self._local_sock = None
        self._remote_sock = None
        self._create_sockets()

    def _create_sockets(self):
        sockets = []
        for addr in (self._local_addr, self._remote_addr):
            addrs = socket.getaddrinfo(addr[0], addr[1], 0,
                                       socket.SOCK_DGRAM, socket.SOL_UDP)
            if len(addrs) == 0:
                raise Exception("can't get addrinfo for %s:%d" % addr)
            af, socktype, proto, canonname, sa = addrs[0]
            sock = socket.socket(af, socktype, proto)
            sock.setblocking(False)
            sockets.append(sock)
        self._local_sock, self._remote_sock = sockets
        self._local_sock.bind(self._local_addr)

    def _rebuild_sockets(self):
        self._id_to_addr.clear()
        self._loop.remove(self._local_sock)
        self._loop.remove(self._remote_sock)
        self._local_sock.close()
        self._remote_sock.close()
        self._create_sockets()

    def add_to_loop(self, loop):
        DNSRelay.add_to_loop(self, loop)

        loop.add(self._local_sock, eventloop.POLL_IN)
        loop.add(self._remote_sock, eventloop.POLL_IN)

    def _handle_local(self, sock):
        try:
            data, addr = sock.recvfrom(BUF_SIZE)
        except (OSError, IOError) as e:
            logging.error(e)
            if eventloop.errno_from_exception(e) == errno.ECONNRESET:
                # just for Windows lol
                self._rebuild_sockets()
            return
        header = asyncdns.parse_header(data)
        if header:
            try:
                req_id = header[0]
                req = asyncdns.parse_response(data)
                self._id_to_addr[req_id] = addr
                self._remote_sock.sendto(data, self._remote_addr)
                logging.info('request %s', req.hostname)
            except Exception as e:
                import traceback

                traceback.print_exc()
                logging.error(e)

    def _handle_remote(self, sock):
        try:
            data, addr = sock.recvfrom(BUF_SIZE)
        except (OSError, IOError) as e:
            logging.error(e)
            if eventloop.errno_from_exception(e) == errno.ECONNRESET:
                # just for Windows lol
                self._rebuild_sockets()
            return
        if data:
            try:
                header = asyncdns.parse_header(data)
                if header:
                    req_id = header[0]
                    res = asyncdns.parse_response(data)
                    logging.info('response %s', res)
                    addr = self._id_to_addr.get(req_id, None)
                    if addr:
                        for answer in res.answers:
                            if answer and answer[0] in GFW_LIST:
                                return
                        self._local_sock.sendto(data, addr)
                        del self._id_to_addr[req_id]
            except Exception as e:
                import traceback
                traceback.print_exc()
                logging.error(e)
                if eventloop.errno_from_exception(e) == errno.EACCES:
                    # when we have changed our ip
                    self._rebuild_sockets()

    def handle_events(self, events):
        for sock, fd, event in events:
            if sock == self._local_sock:
                self._handle_local(sock)
            elif sock == self._remote_sock:
                self._handle_remote(sock)
        now = time.time()
        if now - self._last_time > CACHE_TIMEOUT / 2:
            self._id_to_addr.sweep()

class TCPDNSRelay(DNSRelay):
    def __init__(self, config):
        DNSRelay.__init__(self, config)

        self._local_to_remote = {}
        self._remote_to_local = {}

        addrs = socket.getaddrinfo(self._local_addr[0], self._local_addr[1], 0,
                                   socket.SOCK_STREAM, socket.SOL_TCP)
        if len(addrs) == 0:
            raise Exception("can't get addrinfo for %s:%d" % self._local_addr)
        af, socktype, proto, canonname, sa = addrs[0]
        self._listen_sock = socket.socket(af, socktype, proto)
        self._listen_sock.setblocking(False)
        self._listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._listen_sock.bind(self._local_addr)
        self._listen_sock.listen(1024)

    def _handle_conn(self, sock):
        try:
            local, addr = sock.accept()
            addrs = socket.getaddrinfo(self._remote_addr[0],
                                       self._remote_addr[1], 0,
                                       socket.SOCK_STREAM, socket.SOL_TCP)
            if len(addrs) == 0:
                raise Exception("can't get addrinfo for %s:%d" %
                                self._remote_addr)
            af, socktype, proto, canonname, sa = addrs[0]
            remote = socket.socket(af, socktype, proto)
            local.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
            remote.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
            self._local_to_remote[local] = remote
            self._remote_to_local[remote] = local

            self._loop.add(local, 0)
            self._loop.add(remote, eventloop.POLL_OUT)
            try:
                remote.connect(self._remote_addr)
            except (OSError, IOError) as e:
                if eventloop.errno_from_exception(e) in (errno.EINPROGRESS,
                                                         errno.EAGAIN):
                    pass
                else:
                    raise
        except (OSError, IOError) as e:
            logging.error(e)

    def _destroy(self, local, remote):
        if local in self._local_to_remote:
            self._loop.remove(local)
            self._loop.remove(remote)
            del self._local_to_remote[local]
            del self._remote_to_local[remote]
            local.close()
            remote.close()
        else:
            logging.error('already destroyed')

    def _handle_local(self, local, event):
        remote = self._local_to_remote[local]
        if event & (eventloop.POLL_ERR | eventloop.POLL_HUP):
            self._destroy(local, remote)
        elif event & eventloop.POLL_IN:
            try:
                data = local.recv(BUF_SIZE)
                if not data:
                    self._destroy(local, remote)
                else:
                    remote.send(data)
            except (OSError, IOError) as e:
                self._destroy(local, self._local_to_remote[local])
                logging.error(e)

    def _handle_remote(self, remote, event):
        local = self._remote_to_local[remote]
        if event & (eventloop.POLL_ERR | eventloop.POLL_HUP):
            self._destroy(local, remote)
        elif event & eventloop.POLL_OUT:
            self._loop.modify(remote, eventloop.POLL_IN)
            self._loop.modify(local, eventloop.POLL_IN)
        elif event & eventloop.POLL_IN:
            try:
                data = remote.recv(BUF_SIZE)
                if not data:
                    self._destroy(local, remote)
                else:
                    try:
                        res = asyncdns.parse_response(data[2:])
                        if res:
                            logging.info('response %s', res)
                    except Exception as e:
                        logging.error(e)
                    local.send(data)
            except (OSError, IOError) as e:
                self._destroy(local, remote)
                logging.error(e)

    def add_to_loop(self, loop):
        DNSRelay.add_to_loop(self, loop)
        loop.add(self._listen_sock, eventloop.POLL_IN)

    def handle_events(self, events):
        for sock, fd, event in events:
            if sock == self._listen_sock:
                self._handle_conn(sock)
            elif sock in self._local_to_remote:
                self._handle_local(sock, event)
            elif sock in self._remote_to_local:
                self._handle_remote(sock, event)
                # TODO implement timeout

def main():
    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s %(levelname)-8s %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S', filemode='a+')

    parser = argparse.ArgumentParser(description='Forward DNS requests.')
    parser.add_argument('-b', '--local_address', metavar='BIND_ADDR', type=str,
                        help='address that listens, default: 127.0.0.1',
                        default='127.0.0.1')
    parser.add_argument('-s', '--dns', metavar='DNS', type=str,
                        help='DNS server to use, default: 8.8.8.8',
                        default='8.8.8.8')

    config = vars(parser.parse_args())

    logging.info("starting dns at %s:%d" % (config['local_address'], 53))

    loop = eventloop.EventLoop()

    try:
        udprelay = UDPDNSRelay(config)
        udprelay.add_to_loop(loop)
        tcprelay = TCPDNSRelay(config)
        tcprelay.add_to_loop(loop)
        loop.run()
    except (OSError, IOError) as e:
        logging.error(e)
        if eventloop.errno_from_exception(e) == errno.EACCES:
            logging.info('please use sudo to run this program')
        sys.exit(1)

if __name__ == '__main__':
    main()
<commit_msg>fix windows<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

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

import sys
import time
import socket
import errno
import logging

info = sys.version_info
if not (info[0] == 2 and info[1] >= 7):
    print 'Python 2.7 required'
    sys.exit(1)

import argparse
from shadowsocks import eventloop, asyncdns, lru_cache

BUF_SIZE = 16384

CACHE_TIMEOUT = 10

GFW_LIST = set(["74.125.127.102", "74.125.155.102", "74.125.39.102",
                "74.125.39.113", "209.85.229.138", "128.121.126.139",
                "159.106.121.75", "169.132.13.103", "192.67.198.6",
                "202.106.1.2", "202.181.7.85", "203.161.230.171",
                "203.98.7.65", "207.12.88.98", "208.56.31.43",
                "209.145.54.50", "209.220.30.174", "209.36.73.33",
                "211.94.66.147", "213.169.251.35", "216.221.188.182",
                "216.234.179.13", "243.185.187.39", "37.61.54.158",
                "4.36.66.178", "46.82.174.68", "59.24.3.173", "64.33.88.161",
                "64.33.99.47", "64.66.163.251", "65.104.202.252",
                "65.160.219.113", "66.45.252.237", "72.14.205.104",
                "72.14.205.99", "78.16.49.15", "8.7.198.45", "93.46.8.89"])

class DNSRelay(object):
    def __init__(self, config):
        self._loop = None
        self._config = config
        self._last_time = time.time()

        self._local_addr = (config['local_address'], 53)
        self._remote_addr = (config['dns'], 53)

    def add_to_loop(self, loop):
        if self._loop:
            raise Exception('already add to loop')
        self._loop = loop
        loop.add_handler(self.handle_events)

    def handle_events(self, events):
        pass

class UDPDNSRelay(DNSRelay):
    def __init__(self, config):
        DNSRelay.__init__(self, config)

        self._id_to_addr = lru_cache.LRUCache(CACHE_TIMEOUT)
        self._local_sock = None
        self._remote_sock = None
        self._create_sockets()

    def _create_sockets(self):
        sockets = []
        for addr in (self._local_addr, self._remote_addr):
            addrs = socket.getaddrinfo(addr[0], addr[1], 0,
                                       socket.SOCK_DGRAM, socket.SOL_UDP)
            if len(addrs) == 0:
                raise Exception("can't get addrinfo for %s:%d" % addr)
            af, socktype, proto, canonname, sa = addrs[0]
            sock = socket.socket(af, socktype, proto)
            sock.setblocking(False)
            sockets.append(sock)
        self._local_sock, self._remote_sock = sockets
        self._local_sock.bind(self._local_addr)

    def _rebuild_sockets(self):
        self._id_to_addr.clear()
        self._loop.remove(self._local_sock)
        self._loop.remove(self._remote_sock)
        self._local_sock.close()
        self._remote_sock.close()
        self._create_sockets()
        self._loop.add(self._local_sock, eventloop.POLL_IN)
        self._loop.add(self._remote_sock, eventloop.POLL_IN)

    def add_to_loop(self, loop):
        DNSRelay.add_to_loop(self, loop)

        loop.add(self._local_sock, eventloop.POLL_IN)
        loop.add(self._remote_sock, eventloop.POLL_IN)

    def _handle_local(self, sock):
        try:
            data, addr = sock.recvfrom(BUF_SIZE)
        except (OSError, IOError) as e:
            logging.error(e)
            if eventloop.errno_from_exception(e) == errno.ECONNRESET:
                # just for Windows lol
                self._rebuild_sockets()
            return
        header = asyncdns.parse_header(data)
        if header:
            try:
                req_id = header[0]
                req = asyncdns.parse_response(data)
                self._id_to_addr[req_id] = addr
                self._remote_sock.sendto(data, self._remote_addr)
                logging.info('request %s', req.hostname)
            except Exception as e:
                import traceback

                traceback.print_exc()
                logging.error(e)

    def _handle_remote(self, sock):
        try:
            data, addr = sock.recvfrom(BUF_SIZE)
        except (OSError, IOError) as e:
            logging.error(e)
            if eventloop.errno_from_exception(e) == errno.ECONNRESET:
                # just for Windows lol
                self._rebuild_sockets()
            return
        if data:
            try:
                header = asyncdns.parse_header(data)
                if header:
                    req_id = header[0]
                    res = asyncdns.parse_response(data)
                    logging.info('response %s', res)
                    addr = self._id_to_addr.get(req_id, None)
                    if addr:
                        for answer in res.answers:
                            if answer and answer[0] in GFW_LIST:
                                return
                        self._local_sock.sendto(data, addr)
                        del self._id_to_addr[req_id]
            except Exception as e:
                import traceback
                traceback.print_exc()
                logging.error(e)
                if eventloop.errno_from_exception(e) == errno.EACCES:
                    # when we have changed our ip
                    self._rebuild_sockets()

    def handle_events(self, events):
        for sock, fd, event in events:
            if sock == self._local_sock:
                self._handle_local(sock)
            elif sock == self._remote_sock:
                self._handle_remote(sock)
        now = time.time()
        if now - self._last_time > CACHE_TIMEOUT / 2:
            self._id_to_addr.sweep()

class TCPDNSRelay(DNSRelay):
    def __init__(self, config):
        DNSRelay.__init__(self, config)

        self._local_to_remote = {}
        self._remote_to_local = {}

        addrs = socket.getaddrinfo(self._local_addr[0], self._local_addr[1], 0,
                                   socket.SOCK_STREAM, socket.SOL_TCP)
        if len(addrs) == 0:
            raise Exception("can't get addrinfo for %s:%d" % self._local_addr)
        af, socktype, proto, canonname, sa = addrs[0]
        self._listen_sock = socket.socket(af, socktype, proto)
        self._listen_sock.setblocking(False)
        self._listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._listen_sock.bind(self._local_addr)
        self._listen_sock.listen(1024)

    def _handle_conn(self, sock):
        try:
            local, addr = sock.accept()
            addrs = socket.getaddrinfo(self._remote_addr[0],
                                       self._remote_addr[1], 0,
                                       socket.SOCK_STREAM, socket.SOL_TCP)
            if len(addrs) == 0:
                raise Exception("can't get addrinfo for %s:%d" %
                                self._remote_addr)
            af, socktype, proto, canonname, sa = addrs[0]
            remote = socket.socket(af, socktype, proto)
            local.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
            remote.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
            self._local_to_remote[local] = remote
            self._remote_to_local[remote] = local

            self._loop.add(local, 0)
            self._loop.add(remote, eventloop.POLL_OUT)
            try:
                remote.connect(self._remote_addr)
            except (OSError, IOError) as e:
                if eventloop.errno_from_exception(e) in (errno.EINPROGRESS,
                                                         errno.EAGAIN):
                    pass
                else:
                    raise
        except (OSError, IOError) as e:
            logging.error(e)

    def _destroy(self, local, remote):
        if local in self._local_to_remote:
            self._loop.remove(local)
            self._loop.remove(remote)
            del self._local_to_remote[local]
            del self._remote_to_local[remote]
            local.close()
            remote.close()
        else:
            logging.error('already destroyed')

    def _handle_local(self, local, event):
        remote = self._local_to_remote[local]
        if event & (eventloop.POLL_ERR | eventloop.POLL_HUP):
            self._destroy(local, remote)
        elif event & eventloop.POLL_IN:
            try:
                data = local.recv(BUF_SIZE)
                if not data:
                    self._destroy(local, remote)
                else:
                    remote.send(data)
            except (OSError, IOError) as e:
                self._destroy(local, self._local_to_remote[local])
                logging.error(e)

    def _handle_remote(self, remote, event):
        local = self._remote_to_local[remote]
        if event & (eventloop.POLL_ERR | eventloop.POLL_HUP):
            self._destroy(local, remote)
        elif event & eventloop.POLL_OUT:
            self._loop.modify(remote, eventloop.POLL_IN)
            self._loop.modify(local, eventloop.POLL_IN)
        elif event & eventloop.POLL_IN:
            try:
                data = remote.recv(BUF_SIZE)
                if not data:
                    self._destroy(local, remote)
                else:
                    try:
                        res = asyncdns.parse_response(data[2:])
                        if res:
                            logging.info('response %s', res)
                    except Exception as e:
                        logging.error(e)
                    local.send(data)
            except (OSError, IOError) as e:
                self._destroy(local, remote)
                logging.error(e)

    def add_to_loop(self, loop):
        DNSRelay.add_to_loop(self, loop)
        loop.add(self._listen_sock, eventloop.POLL_IN)

    def handle_events(self, events):
        for sock, fd, event in events:
            if sock == self._listen_sock:
                self._handle_conn(sock)
            elif sock in self._local_to_remote:
                self._handle_local(sock, event)
            elif sock in self._remote_to_local:
                self._handle_remote(sock, event)
                # TODO implement timeout

def main():
    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s %(levelname)-8s %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S', filemode='a+')

    parser = argparse.ArgumentParser(description='Forward DNS requests.')
    parser.add_argument('-b', '--local_address', metavar='BIND_ADDR', type=str,
                        help='address that listens, default: 127.0.0.1',
                        default='127.0.0.1')
    parser.add_argument('-s', '--dns', metavar='DNS', type=str,
                        help='DNS server to use, default: 8.8.8.8',
                        default='8.8.8.8')

    config = vars(parser.parse_args())

    logging.info("starting dns at %s:%d" % (config['local_address'], 53))

    loop = eventloop.EventLoop()

    try:
        udprelay = UDPDNSRelay(config)
        udprelay.add_to_loop(loop)
        tcprelay = TCPDNSRelay(config)
        tcprelay.add_to_loop(loop)
        loop.run()
    except (OSError, IOError) as e:
        logging.error(e)
        if eventloop.errno_from_exception(e) == errno.EACCES:
            logging.info('please use sudo to run this program')
        sys.exit(1)

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


<commit_before>"""
Custom storage for django with Mosso Cloud Files backend.
Created by Rich Leland <rich@richleland.com>.
"""
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.core.files import File
from django.core.files.storage import Storage
from django.utils.text import get_valid_filename

try:
    import cloudfiles
    from cloudfiles.errors import NoSuchObject
except ImportError:
    raise ImproperlyConfigured("Could not load cloudfiles dependency. See "
                               "http://www.mosso.com/cloudfiles.jsp.")

# TODO: implement TTL into cloudfiles methods
CLOUDFILES_TTL = getattr(settings, 'CLOUDFILES_TTL', 600)

def cloudfiles_upload_to(self, filename):
    """
    Simple, custom upload_to because Cloud Files doesn't support
    nested containers (directories).

    Actually found this out from @minter:
    @richleland The Cloud Files APIs do support pseudo-subdirectories, by
    creating zero-byte files with type application/directory.

    May implement in a future version.
    """
    return get_valid_filename(filename)

class CloudFilesStorage(Storage):
    """
    Custom storage for Mosso Cloud Files.
    """
    default_quick_listdir = True

    def __init__(self, username=None, api_key=None, container=None,
                 connection_kwargs=None):
        """
        Initialize the settings for the connection and container.
        """
        self.username = username or settings.CLOUDFILES_USERNAME
        self.api_key = api_key or settings.CLOUDFILES_API_KEY
        self.container_name = container or settings.CLOUDFILES_CONTAINER
        self.connection_kwargs = connection_kwargs or {}

    def __getstate__(self):
        """
        Return a picklable representation of the storage.
        """
        return dict(username=self.username,
                    api_key=self.api_key,
                    container_name=self.container_name,
                    connection_kwargs=self.connection_kwargs)

    def _get_connection(self):
        if not hasattr(self, '_connection'):
            self._connection = cloudfiles.get_connection(self.username,
                                    self.api_key, **self.connection_kwargs)
        return self._connection

    def _set_connection(self, value):
        self._connection = value

    connection = property(_get_connection, _set_connection)

    def _get_container(self):
        if not hasattr(self, '_container'):
            self.container = self.connection.get_container(
                                                        self.container_name)
        return self._container

    def _set_container(self, container):
        """
        Set the container, making it publicly available (on Limelight CDN) if
        it is not already.
        """
        if not container.is_public():
            container.make_public()
        if hasattr(self, '_container_public_uri'):
            delattr(self, '_container_public_uri')
        self._container = container

    container = property(_get_container, _set_container)

    def _get_container_url(self):
        if not hasattr(self, '_container_public_uri'):
            self._container_public_uri = self.container.public_uri()
        return self._container_public_uri

    container_url = property(_get_container_url)

    def _get_cloud_obj(self, name):
        """
        Helper function to get retrieve the requested Cloud Files Object.
        """
        return self.container.get_object(name)

    def _open(self, name, mode='rb'):
        """
        Return the CloudFilesStorageFile.
        """
        return CloudFilesStorageFile(storage=self, name=name)

    def _save(self, name, content):
        """
        Use the Cloud Files service to write ``content`` to a remote file
        (called ``name``).
        """

        cloud_obj = self.container.create_object(name)
        cloud_obj.size = content.size

        content.open()
        # If the content type is available, pass it in directly rather than
        # getting the cloud object to try to guess.
        if hasattr(content.file, 'content_type'):
            cloud_obj.content_type = content.file.content_type
        cloud_obj.send(content)
        content.close()
        return name

    def delete(self, name):
        """
        Deletes the specified file from the storage system.
        """
        # If the file exists, delete it.
        if self.exists(name):
            self.container.delete_object(name)

    def exists(self, name):
        """
        Returns True if a file referenced by the given name already exists in
        the storage system, or False if the name is available for a new file.
        """
        try:
            self._get_cloud_obj(name)
            return True
        except NoSuchObject:
            return False

    def listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple; the
        first being an empty list of directories (not available for quick-
        listing), the second being a list of filenames.

        If the list of directories is required, use the full_listdir method.
        """
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(path=path):
            files.append(name[path_len:])
        return ([], files)

    def full_listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple of lists;
        the first item being directories, the second item being files.

        On large containers, this may be a slow operation for root containers
        because every single object must be returned (cloudfiles does not
        provide an explicit way of listing directories).
        """
        dirs = set()
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(prefix=path):
            name = name[path_len:]
            slash = name[1:-1].find('/') + 1
            if slash:
                dirs.add(name[:slash])
            elif name:
                files.append(name)
        dirs = list(dirs)
        dirs.sort()
        return (dirs, files)

    def size(self, name):
        """
        Returns the total size, in bytes, of the file specified by name.
        """
        return self._get_cloud_obj(name).size

    def url(self, name):
        """
        Returns an absolute URL where the file's contents can be accessed
        directly by a web browser.
        """
        return '%s/%s' % (self.container_url, name)

class CloudFilesStorageFile(File):
    closed = False

    def __init__(self, storage, name, *args, **kwargs):
        self._storage = storage
        super(CloudFilesStorageFile, self).__init__(file=None, name=name,
                                                    *args, **kwargs)

    def _get_size(self):
        if not hasattr(self, '_size'):
            self._size = self._storage.size(self.name)
        return self._size

    def _set_size(self, size):
        self._size = size

    size = property(_get_size, _set_size)

    def _get_file(self):
        if not hasattr(self, '_file'):
            self._file = self._storage._get_cloud_obj(self.name)
        return self._file

    def _set_file(self, value):
        if value is None:
            if hasattr(self, '_file'):
                del self._file
        else:
            self._file = value

    file = property(_get_file, _set_file)

    def read(self, num_bytes=None):
        data = self.file.read(size=num_bytes or -1, offset=self._pos)
        self._pos += len(data)
        return data

    def open(self, *args, **kwargs):
        """
        Open the cloud file object.
        """
        self.file
        self._pos = 0

    def close(self, *args, **kwargs):
        self._pos = 0

    @property
    def closed(self):
        return not hasattr(self, '_file')

    def seek(self, pos):
        self._pos = pos
<commit_msg>Added the option to set the connection kwargs inside the django settings file<commit_after>"""
Custom storage for django with Mosso Cloud Files backend.
Created by Rich Leland <rich@richleland.com>.
"""
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.core.files import File
from django.core.files.storage import Storage
from django.utils.text import get_valid_filename

try:
    import cloudfiles
    from cloudfiles.errors import NoSuchObject
except ImportError:
    raise ImproperlyConfigured("Could not load cloudfiles dependency. See "
                               "http://www.mosso.com/cloudfiles.jsp.")

# TODO: implement TTL into cloudfiles methods
CLOUDFILES_TTL = getattr(settings, 'CLOUDFILES_TTL', 600)

def cloudfiles_upload_to(self, filename):
    """
    Simple, custom upload_to because Cloud Files doesn't support
    nested containers (directories).

    Actually found this out from @minter:
    @richleland The Cloud Files APIs do support pseudo-subdirectories, by
    creating zero-byte files with type application/directory.

    May implement in a future version.
    """
    return get_valid_filename(filename)

class CloudFilesStorage(Storage):
    """
    Custom storage for Mosso Cloud Files.
    """
    default_quick_listdir = True

    def __init__(self, username=None, api_key=None, container=None,
                 connection_kwargs=None):
        """
        Initialize the settings for the connection and container.
        """
        self.username = username or settings.CLOUDFILES_USERNAME
        self.api_key = api_key or settings.CLOUDFILES_API_KEY
        self.container_name = container or settings.CLOUDFILES_CONTAINER
        self.connection_kwargs = connection_kwargs or settings.CLOUDFILES_CONNECTION_KWARGS or {}

    def __getstate__(self):
        """
        Return a picklable representation of the storage.
        """
        return dict(username=self.username,
                    api_key=self.api_key,
                    container_name=self.container_name,
                    connection_kwargs=self.connection_kwargs)

    def _get_connection(self):
        if not hasattr(self, '_connection'):
            self._connection = cloudfiles.get_connection(self.username,
                                    self.api_key, **self.connection_kwargs)
        return self._connection

    def _set_connection(self, value):
        self._connection = value

    connection = property(_get_connection, _set_connection)

    def _get_container(self):
        if not hasattr(self, '_container'):
            self.container = self.connection.get_container(
                                                        self.container_name)
        return self._container

    def _set_container(self, container):
        """
        Set the container, making it publicly available (on Limelight CDN) if
        it is not already.
        """
        if not container.is_public():
            container.make_public()
        if hasattr(self, '_container_public_uri'):
            delattr(self, '_container_public_uri')
        self._container = container

    container = property(_get_container, _set_container)

    def _get_container_url(self):
        if not hasattr(self, '_container_public_uri'):
            self._container_public_uri = self.container.public_uri()
        return self._container_public_uri

    container_url = property(_get_container_url)

    def _get_cloud_obj(self, name):
        """
        Helper function to get retrieve the requested Cloud Files Object.
        """
        return self.container.get_object(name)

    def _open(self, name, mode='rb'):
        """
        Return the CloudFilesStorageFile.
        """
        return CloudFilesStorageFile(storage=self, name=name)

    def _save(self, name, content):
        """
        Use the Cloud Files service to write ``content`` to a remote file
        (called ``name``).
        """

        cloud_obj = self.container.create_object(name)
        cloud_obj.size = content.size

        content.open()
        # If the content type is available, pass it in directly rather than
        # getting the cloud object to try to guess.
        if hasattr(content.file, 'content_type'):
            cloud_obj.content_type = content.file.content_type
        cloud_obj.send(content)
        content.close()
        return name

    def delete(self, name):
        """
        Deletes the specified file from the storage system.
        """
        # If the file exists, delete it.
        if self.exists(name):
            self.container.delete_object(name)

    def exists(self, name):
        """
        Returns True if a file referenced by the given name already exists in
        the storage system, or False if the name is available for a new file.
        """
        try:
            self._get_cloud_obj(name)
            return True
        except NoSuchObject:
            return False

    def listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple; the
        first being an empty list of directories (not available for quick-
        listing), the second being a list of filenames.

        If the list of directories is required, use the full_listdir method.
        """
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(path=path):
            files.append(name[path_len:])
        return ([], files)

    def full_listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple of lists;
        the first item being directories, the second item being files.

        On large containers, this may be a slow operation for root containers
        because every single object must be returned (cloudfiles does not
        provide an explicit way of listing directories).
        """
        dirs = set()
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(prefix=path):
            name = name[path_len:]
            slash = name[1:-1].find('/') + 1
            if slash:
                dirs.add(name[:slash])
            elif name:
                files.append(name)
        dirs = list(dirs)
        dirs.sort()
        return (dirs, files)

    def size(self, name):
        """
        Returns the total size, in bytes, of the file specified by name.
        """
        return self._get_cloud_obj(name).size

    def url(self, name):
        """
        Returns an absolute URL where the file's contents can be accessed
        directly by a web browser.
        """
        return '%s/%s' % (self.container_url, name)

class CloudFilesStorageFile(File):
    closed = False

    def __init__(self, storage, name, *args, **kwargs):
        self._storage = storage
        super(CloudFilesStorageFile, self).__init__(file=None, name=name,
                                                    *args, **kwargs)

    def _get_size(self):
        if not hasattr(self, '_size'):
            self._size = self._storage.size(self.name)
        return self._size

    def _set_size(self, size):
        self._size = size

    size = property(_get_size, _set_size)

    def _get_file(self):
        if not hasattr(self, '_file'):
            self._file = self._storage._get_cloud_obj(self.name)
        return self._file

    def _set_file(self, value):
        if value is None:
            if hasattr(self, '_file'):
                del self._file
        else:
            self._file = value

    file = property(_get_file, _set_file)

    def read(self, num_bytes=None):
        data = self.file.read(size=num_bytes or -1, offset=self._pos)
        self._pos += len(data)
        return data

    def open(self, *args, **kwargs):
        """
        Open the cloud file object.
        """
        self.file
        self._pos = 0

    def close(self, *args, **kwargs):
        self._pos = 0

    @property
    def closed(self):
        return not hasattr(self, '_file')

    def seek(self, pos):
        self._pos = pos
<|endoftext|>


<commit_before>import argparse

__author__ = 'moskupols'

def evolution(title=None, resource=None, *, modifiers: tuple=None, explanations=None):
    """
    Generates evolution stories of given explanations or explanations with known title
    being affected by modifiers of given resource or just given modifiers.

    Each story is a tuple of pairs (modifier, result). The first element of this tuple is always
     (None, initial explanation).

    For call you should specify modifiers and explanations, but you can do it implicitly
    by specifying resource and title to seek in the resource's output.

    If explanations is specified and is instance of Explanation, the function returns evolution story.
    Otherwise it returns sequence (start_explanation, evolution story).

    :param title: title of explanation to seek in exactly the same format as generated by the resource.
    :param resource: name or resource itself.
    :param modifiers: if not specified, modifiers are taken from resource.modifiers.
    :param explanations: if not specified, explanations are taken from (e for e in resource if e.title == title).
    :return:
    """
    from copy import copy
    from hb_res.explanations import Explanation
    from preparation.resources.Resource import resource_by_name

    def evolute(start, mods):
        cur = copy(start)
        yield (None, cur)
        for m in mods:
            cur = m(cur)
            yield (m, cur)
            if cur is None:
                break

    if isinstance(resource, str):
        resource = resource_by_name(resource)()

    assert modifiers is not None or hasattr(resource, 'modifiers')
    if modifiers is None:
        modifiers = resource.modifiers

    assert (title is not None and resource is not None) ^ (explanations is not None)
    if explanations is None:
        explanations = (e for e in resource if e.title == title)

    if isinstance(explanations, Explanation):
        return tuple(evolute(explanations, modifiers))
    else:
        for start_expl in explanations:
            yield (start_expl, tuple(evolute(start_expl, modifiers)))

def make_argparser():
    from preparation.resources.Resource import names_registered

    parser = argparse.ArgumentParser(description='View how some explanation(s) evolute ')

    formats = {
        'oneline': '{mod!s:>50} -> ("{result.title}" := "{result.text}")',
        'wide': '{mod!s:>60} -> {result!r}',
        'twoline': 'after {mod}:\n --> {result!r}'
    }

    fmt_args = parser.add_mutually_exclusive_group()
    fmt_args.add_argument('--format',
                          default=formats['twoline'],
                          help='''Self argument for str.format call while printing evolution.
                          If omitted, same as --twoline.
                          ''')
    for name, fmt in formats.items():
        fmt_args.add_argument('--' + name,
                              dest='format',
                              action='store_const',
                              const=fmt,
                              help='Format as ' + repr(fmt))

    parser.add_argument('resource',
                        choices=names_registered(),
                        help='resource to use')
    parser.add_argument('title',
                        nargs='+',
                        help='''title (titles) to show.
                        They will be used to find raw explanations yielded by the resource's parser,
                        so should be exactly in the same format the parser gives it.''')
    return parser

def main(args=None):
    if not isinstance(args, argparse.Namespace):
        parser = make_argparser()
        args = parser.parse_args(args)
    first_story = True
    for title in args.title:
        stories = evolution(title=title, resource=args.resource)
        for start_expl, story in stories:
            if first_story:
                first_story = False
            else:
                print()
            for mod, result in story:
                print(args.format.format(mod=mod, result=result))

if __name__ == '__main__':
    main()
<commit_msg>Replace resource specification with trunk specification<commit_after>import argparse

__author__ = 'moskupols'

def evolution(title=None, resource=None, *, modifiers: tuple=None, explanations=None):
    """
    Generates evolution stories of given explanations or explanations with known title
    being affected by modifiers of given resource or just given modifiers.

    Each story is a tuple of pairs (modifier, result). The first element of this tuple is always
     (None, initial explanation).

    For call you should specify modifiers and explanations, but you can do it implicitly
    by specifying resource and title to seek in the resource's output.

    If explanations is specified and is instance of Explanation, the function returns evolution story.
    Otherwise it returns sequence (start_explanation, evolution story).

    :param title: title of explanation to seek in exactly the same format as generated by the resource.
    :param resource: name or resource itself.
    :param modifiers: if not specified, modifiers are taken from resource.modifiers.
    :param explanations: if not specified, explanations are taken from (e for e in resource if e.title == title).
    :return:
    """
    from copy import copy
    from hb_res.explanations import Explanation
    from preparation.resources.Resource import resource_by_name

    def evolute(start, mods):
        cur = copy(start)
        yield (None, cur)
        for m in mods:
            cur = m(cur)
            yield (m, cur)
            if cur is None:
                break

    if isinstance(resource, str):
        resource = resource_by_name(resource)()

    assert modifiers is not None or hasattr(resource, 'modifiers')
    if modifiers is None:
        modifiers = resource.modifiers

    assert (title is not None and resource is not None) ^ (explanations is not None)
    if explanations is None:
        explanations = (e for e in resource if e.title == title)

    if isinstance(explanations, Explanation):
        return tuple(evolute(explanations, modifiers))
    else:
        for start_expl in explanations:
            yield (start_expl, tuple(evolute(start_expl, modifiers)))

def make_argparser():
    from preparation.resources.Resource import names_registered

    parser = argparse.ArgumentParser(description='View how some explanation(s) evolute ')

    formats = {
        'oneline': '{mod!s:>50} -> ("{result.title}" := "{result.text}")',
        'wide': '{mod!s:>60} -> {result!r}',
        'twoline': 'after {mod}:\n --> {result!r}'
    }
    trunks = [name.replace('Resource', '') for name in names_registered()]

    fmt_args = parser.add_mutually_exclusive_group()
    fmt_args.add_argument('--format',
                          default=formats['twoline'],
                          help='''Self argument for str.format call while printing evolution.
                          If omitted, same as --twoline.
                          ''')
    for name, fmt in formats.items():
        fmt_args.add_argument('--' + name,
                              dest='format',
                              action='store_const',
                              const=fmt,
                              help='Format as ' + repr(fmt))

    parser.add_argument('resource',
                        choices=trunks,
                        help='resource to use')
    parser.add_argument('title',
                        nargs='+',
                        help='''title (titles) to show.
                        They will be used to find raw explanations yielded by the resource's parser,
                        so should be exactly in the same format the parser gives it.''')
    return parser

def main(args=None):
    if not isinstance(args, argparse.Namespace):
        parser = make_argparser()
        args = parser.parse_args(args)
    first_story = True
    for title in args.title:
        stories = evolution(title=title, resource=args.resource + 'Resource')
        for start_expl, story in stories:
            if first_story:
                first_story = False
            else:
                print()
            for mod, result in story:
                print(args.format.format(mod=mod, result=result))

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