# !/usr/bin/env python
# encoding: utf-8

import os
from tqdm import tqdm
import logging
import random
import numpy as np
import torch
from tqdm import trange
import torch.nn as nn
from datetime import datetime
from utils import to_np, trim_seqs, trim_seqs_beam, get_classify_types, get_func_id_list
from torch_geometric.loader import DataLoader
from sklearn.metrics import classification_report, top_k_accuracy_score

logger = logging.getLogger(__name__)

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

def get_input_from_batch(batch, args):
    node_pos = batch[0].ptr.to(args.device)
    x = batch[0].x.to(args.device)
    edge_index = batch[0].edge_index.to(args.device)
    labels = batch[1].to(args.device)
    func_ids = batch[2]

    return x, edge_index, node_pos, labels, func_ids

def trainer(args, model, train_dataset, test_dataset, type2idx, idx2type):
    # Training model
    start_time = datetime.now()
    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    eval_dataloader = DataLoader(test_dataset, batch_size=args.test_batch_size, shuffle=False)

    # logger.info('Num examples = %d', len(train_dataset))

    if args.model == 'classify':

        if not args.test_only:

            train_classify(args, model, train_dataloader, eval_dataloader, idx2type)

            end_time = datetime.now()
            logger.info('Model training spends %s' % str((end_time - start_time).seconds))

        # Test model
        base_path = os.path.join(args.model_save_path, args.language, args.module)
        path = base_path + '/return_pred_data_' + args.model + '_' + args.pos + '_' + \
               str(args.batch_size) + '_' + str(args.learning_rate) + '_' + str(args.epochs)
        model.load_state_dict(torch.load(path + '.pt'))

        start_time = datetime.now()

        indicators = classify_eval(args, model, eval_dataloader, idx2type)

        end_time = datetime.now()

        internal_time = int(str((end_time - start_time).seconds))

        return indicators, internal_time

def train_classify(args, model, train_dataset, test_dataset, idx2type):

    parameters = filter(lambda param: param.requires_grad, model.parameters())
    optimizer = torch.optim.Adam(parameters, lr=args.learning_rate)

    logger.info('****** Running training ******')

    global_step = 0

    model.zero_grad()
    train_iterator = trange(int(args.epochs), desc='Epoch')
    set_seed(args)

    total_loss = 0.
    last_loss = 10000000

    best_top1 = 0.

    for cur_epoch in train_iterator:

        logger.info('============================= %s =======================' % str(cur_epoch))

        # step = 0
        for batch in tqdm(train_dataset):

            model.train()

            x, edge_index, node_pos, labels, func_ids = get_input_from_batch(batch, args)

            logit = model(x, edge_index, node_pos)

            # Calculate loss for multi-tags
            loss = get_loss_classify(logit, labels)

            # if args.gradient_accumulation_steps > 1:
            #     loss = loss / args.gradient_accumulation_steps

            loss.backward()

            torch.nn.utils.clip_grad_norm_(
                model.parameters(), args.max_grad_norm)

            total_loss += loss.item()

            optimizer.step()
            model.zero_grad()
            global_step += 1

            # if args.logging_steps > 0 and global_step != 1 and global_step % args.logging_steps == 0:
            if args.logging_steps > 0 and global_step != 1:

                avg_loss = total_loss / global_step

                logger.info('********** Training loss %s ***************', str(avg_loss))

                if avg_loss < last_loss:
                    # save model checkpoint according to the loss
                    base_path = os.path.join(args.model_save_path, args.language, args.module)
                    if not os.path.exists(base_path):
                        os.makedirs(base_path)
                    path = base_path + '/return_pred_data_' + args.model + '_' + args.pos + '_' + \
                           str(args.batch_size) + '_' + str(args.learning_rate) + '_' + str(args.epochs)
                    torch.save(model.state_dict(), path + '.pt')

def get_loss_classify(logits, y_true):
    loss = nn.CrossEntropyLoss()
    output = loss(logits, y_true)
    return output

def classify_eval(args, model, eval_dataset, idx2type):
    print('evaluate model')

    # Eval
    logger.info('****** Running evaluation *******')

    eval_loss = 0.0
    nb_eval_steps = 0
    preds = None
    out_label_ids = None
    func_id_lists = None
    for batch in tqdm(eval_dataset):
        model.eval()

        with torch.no_grad():

            x, edge_index, node_pos, labels, func_ids = get_input_from_batch(batch, args)

            logits = model(x, edge_index, node_pos)

            tmp_eval_loss = get_loss_classify(logits, labels)

            eval_loss += tmp_eval_loss.mean().item()

        nb_eval_steps += 1

        if preds is None:
            preds = logits.detach().cpu().numpy()
            out_label_ids = labels.detach().cpu().numpy()
            if args.language == 'solidity':
                func_id_lists = func_ids.detach().cpu().numpy()
        else:
            preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
            out_label_ids = np.append(out_label_ids, labels.detach().cpu().numpy(), axis=0)
            if args.language == 'solidity':
                func_id_lists = np.append(func_id_lists, func_ids.detach().cpu().numpy(), axis=0)

    eval_loss = eval_loss / nb_eval_steps

    logger.info('********** Current loss %s ***************', str(eval_loss))

    indicators = classify_indicator(preds, out_label_ids, func_id_lists, idx2type, args)

    return indicators

def classify_indicator(logits, y_true, func_id_lists, idx2type, args):
    ret = []
    func_id_pred_trues = []

    y_pred = np.argmax(logits, axis=1)

    top1_acc = top_k_accuracy_score(y_true, logits, k=1, labels=range(0, args.num_classes))
    top3_acc = top_k_accuracy_score(y_true, logits, k=3, labels=range(0, args.num_classes))
    top5_acc = top_k_accuracy_score(y_true, logits, k=5, labels=range(0, args.num_classes))

    ret.append({
        'top1_acc': str(top1_acc),
        'top3_acc': str(top3_acc),
        'top5_acc': str(top5_acc),
    })

    return ret


import dotenv
import lucene
import streamlit as st
from pybool_ir.experiments.retrieval import LuceneSearcher
from streamlit.components.v1 import html

from config import AppConfig

assert lucene.getVMEnv() or lucene.initVM()
jvm = lucene.getVMEnv()
jvm.attachCurrentThread()

# Config ----------------------------------------------------------------------
config = AppConfig(dotenv.dotenv_values(".env")["PBIRUI_CONFIG_FILE"])
parser = config.parser
indexer = config.indexer
default_query = config.default_query
page_title = config.page_title
render_serp_item = config.render_serp_item
# -----------------------------------------------------------------------------

if "page" not in st.session_state:
    st.session_state.page = 0
if "serp_rendered" not in st.session_state:
    st.session_state.serp_rendered = False
results_per_page = 10

def util_search():
    assert lucene.getVMEnv() or lucene.initVM()
    jvm = lucene.getVMEnv()
    jvm.attachCurrentThread()

    with LuceneSearcher(indexer) as searcher:
        try:
            query = parser.parse_lucene(inp_query)
        except Exception as e:
            with cnt_results:
                st.error(str(e))
            return
        result_iterator = searcher.index.search(query)
        result_count = searcher.index.count(query)
        return result_iterator, result_count

def search():
    result_iterator, result_count = util_search()
    if st.session_state.page * results_per_page > result_count:
        with cnt_results:
            st.error("no more results")
        st.session_state.page -= 1
        return
    results = [x for x in result_iterator[st.session_state["page"]:st.session_state["page"] + results_per_page]]
    with cnt_results:
        st.write(f"{result_count} results, page {st.session_state['page'] + 1}")
        for result in results:
            render_serp_item(result)

def first_page():
    st.session_state.serp_rendered = True
    st.session_state.page = 0
    search()

def next_page():
    cnt_results.write("")
    st.session_state.page += 1
    search()

def prev_page():
    cnt_results.write("")
    if st.session_state.page > 0:
        st.session_state.page -= 1
    search()

st.title(page_title)

inp_query = st.text_input("Query", default_query)
btn_search = st.button("Search", on_click=first_page)
cnt_results = st.container()
cnt_results.write("")
col_pages1, col_pages2 = st.columns(2)
with col_pages1:
    if st.session_state.page > 0:
        btn_prev = st.button("previous page", on_click=prev_page, use_container_width=True)
with col_pages2:
    if st.session_state.serp_rendered:
        btn_more = st.button("next page", on_click=next_page, use_container_width=True)

# Scroll to top of page when page changes.
html(
    f"""
    <p>{st.session_state.page}</p>
    <script>
        var input = window.parent.document.querySelectorAll("input[type=text]");
        input[0].focus();
    </script>
    """,
    height=0,
)


import json
import logging
import requests
import uuid
from datetime import datetime
from django.shortcuts import render
from django.views.decorators.csrf import csrf_exempt
from rest_framework.response import Response
from rest_framework import status
from rest_framework.decorators import api_view

from django.conf import settings
from .models import User, Policy, Device, Stay, Entity, Consent_Entity, Consent_Device, Receipt, Stay_Receipt, State
from .forms import PolicyForm, DeviceForm, UserForm, StayForm, EntityForm

logger = logging.getLogger(__name__)

''' ##########################################################################
        INTERFACE
##########################################################################'''

'''
    Initial page just to init the demo
'''
def index(request):
    return render(request, 'index.html')

'''
    Render the HTML to show the form to register user
'''
def formRegisterUser(request):
    form = UserForm()
    return render(request, 'registerUser.html', {'form': form})

'''
    Render the HTML to show the form to register stay
'''
def formRegisterStay(request):
    form = StayForm()
    return render(request, 'registerStay.html', {'form': form})

'''
    Render the HTML to show the form to add policy
'''
def formAddPolicy(request):
    form = PolicyForm()
    return render(request, 'addPolicy.html', {'form': form})

'''
    Render the HTML to show the form to add device
'''
def formAddDevice(request):
    form = DeviceForm()
    return render(request, 'addDevice.html', {'form': form})

'''
    Render the HTML to show the form to add entity
'''
def formAddEntity(request):
    form = EntityForm()
    return render(request, 'addEntity.html', {'form': form})

'''
    Render the HTML to show the receipt request
'''
def formReceiptRequest(request):
    form = ReceiptForm()
    return render(request, 'requestReceipt.html', {'form': form})

'''
    List all policies
'''
def listPolicies(request):
    policies = []

    policy_object = Policy.objects.all()
    for p in policy_object:
        di = {'id': p.policyid, 'policy': p.policy}

        policies.append(di)

    return render(request, 'listPolicies.html', {'Policies': policies})

'''
    List all devices
'''
def listDevices(request):
    devices = []

    device_object = Device.objects.all()
    for d in device_object:
        di = {'id': d.deviceid, 'name': d.device, 'policy': d.policyid}

        devices.append(di)

    return render(request, 'listDevices.html', {'Devices': devices})

'''
    List all entity
'''
def listEntities(request):
    entities = []

    entity_object = Entity.objects.all()
    for d in entity_object:
        di = {'id': d.entityid, 'name': d.entity, 'policy': d.policyid}

        entities.append(di)

    return render(request, 'listEntities.html', {'Entities': entities})

'''
    List all users
'''
def listUsers(request):
    users = []

    user_object = User.objects.all()
    for u in user_object:
        di = {'email': u.email, 'FirstName': u.firstname, 'LastName': u.lastname}

        users.append(di)

    return render(request, 'listUsers.html', {'Users': users})

'''
    List temporary receipts - receipts to sign
'''
def listReceipts(request):
    # get email from url parameter
    email = request.GET.get('email', None)

    receipts = []
    emails = []

    # get all users
    users = User.objects.all()
    for u in users:
        emails.append(u.email)

    # get the first email for available users
    if email is None and len(emails) > 0:
        email = emails[0]

    receipt_object = Receipt.objects.filter(stayid__email=email)
    for r in receipt_object:
        print(r)
        ri = {
            'receipt': json.dumps(r.json_receipt),
            'timestampStored': r.timestamp_stored,
            'timestampCreated': r.timestamp_created,
            'stayId': r.stayid.pk,
            'pk': r.id
            }
        receipts.append(ri)
    return render(request, 'listReceipts.html', {'email':email, 'emails':emails, 'receipts': receipts})

'''
    List signed receipts
'''
def listSReceipts(request):
    # get email from url parameter
    email = request.GET.get('email', None)

    receipts = []
    emails = []

    # get all users
    users = User.objects.all()
    for u in users:
        emails.append(u.email)

    # get the first email for available users
    if email is None and len(emails) > 0:
        email = emails[0]

    qs = Stay_Receipt.objects.filter(stayid__email=email)
    for stay in qs:
        url = settings.RECEIPTGET
        p = {'email': email, 'receiptid': stay.receiptid}
        x = requests.get(url, params=p)
        receipt_object = json.loads(x.text)['receipt']
        ri = {'receipt': json.dumps(receipt_object), 'datein': stay.stayid.datein, 'dateout': stay.stayid.dateout, 'stayId': stay.pk}
        receipts.append(ri)

    return render(request, 'listSReceipts.html', {'email':email, 'emails':emails, 'receipts': receipts})

def notifications(request):
    requests = []
    qs = Stay.objects.filter(state=State.re)
    for stay in qs:
        requests.append({'id':stay.pk, 'email':stay.email, 'datein': stay.datein, 'dateout': stay.dateout})
    print(requests)
    return render(request, 'notifications.html', {'requests':requests})

'''##########################################################################
        API
##########################################################################'''
'''
    Register temporary user in the system and post data on data retention
'''
@csrf_exempt
@api_view(('POST',))
def registerUser(request):
    parameters = json.loads(request.body)
    firstname = parameters['firstname']
    lastname = parameters['lastname']
    email = parameters['email']

    try:
        User.objects.create(email=email, firstname=firstname, lastname=lastname)
    except Exception as e:
        print(e)
        return Response('Cannot create the user record', status=status.HTTP_400_BAD_REQUEST)

    return Response(status=status.HTTP_201_CREATED)

@csrf_exempt
@api_view(('POST',))
def registerStay(request):
    parameters = json.loads(request.body)
    print(parameters)
    datein = parameters['datein']
    dateout = parameters['dateout']
    email = parameters['email']
    device_pk = [x[7:] for x in parameters.keys() if x.startswith('device_')]
    entity_pk = [x[7:] for x in parameters.keys() if x.startswith('entity_')]

    version = 1
    language = 'EN'
    selfservicepoint = 'cassiopeia.id'
    userid = email
    devices = []
    entities = []
    otherinfo = ''

    try:
        #Create stay info in db
        user = User.objects.get(email=email)
        stay = Stay.objects.create(datein=datein, dateout=dateout, email=user)
        for pk in device_pk:
            device = Device.objects.get(deviceid = pk)
            consent_d = parameters[f'device_{pk}'] == 'True'
            Consent_Device.objects.create(stayid=stay, deviceid=device, consent=consent_d)
            devices.append({'ID': pk, 'Policy': device.policyid.pk, 'Consent': consent_d})

        for pk in entity_pk:
            entity = Entity.objects.get(entityid = pk)
            consent_e = parameters[f'entity_{pk}'] == 'True'
            Consent_Entity.objects.create(stayid=stay, entityid=entity, consent=consent_e)
            entities.append({'ID': pk, 'Policy': entity.policyid.pk, 'Consent': consent_e})

        url = settings.RECEIPTGENERATION
        r = {
            'version':version,
            'language':language,
            'selfservicepoint':selfservicepoint,
            'userid':userid,
            #'privacyid':privacyid,
            'devices':devices,
            'entities':entities,
            'otherinfo':otherinfo}
        x = requests.get(url, json=r)

        result = json.loads(x.text)

        receipt = result['receipt']
        timestamp = result['timestamp']

        dt_object = datetime.fromtimestamp(timestamp)

        Receipt.objects.create(timestamp_created=dt_object, json_receipt=receipt, stayid=stay)

        #url = settings.DATA_RETENTION_STAY
        #user = {'datein':datein, 'dateout':dateout, 'email':email}
        #x = requests.post(url, data=user)

    except Exception as e:
        print(e)
        return Response('Cannot create the stay record', status=status.HTTP_400_BAD_REQUEST)

    return Response(status=status.HTTP_201_CREATED)

'''
    Create policy and store it on the database to reuse
'''
@csrf_exempt
@api_view(('POST',))
def addPolicy(request):
    parameters = json.loads(request.body)
    policy = parameters['policy']

    try:
        Policy.objects.create(policy=policy)
    except:
        return Response('Cannot create the policy record', status=status.HTTP_400_BAD_REQUEST)

    return Response(status=status.HTTP_201_CREATED)

'''
    Add the available devices
'''
@csrf_exempt
@api_view(('POST',))
def addDevice(request):
    print(f'Add Device into DB')
    parameters = json.loads(request.body)
    print(parameters)
    device = parameters['device']
    policyid = parameters['policyid']

    try:
        policy = Policy.objects.get(policyid=policyid)
        Device.objects.create(device=device, policyid=policy)
    except Exception as e:
        print(e)
        return Response('Cannot create the device record', status=status.HTTP_400_BAD_REQUEST)

    return Response(status=status.HTTP_201_CREATED)

'''
    Add the available entities
'''
@csrf_exempt
@api_view(('POST',))
def addEntity(request):
    print(f'Add Entity into DB')
    parameters = json.loads(request.body)
    print(parameters)
    entity = parameters['entity']
    policyid = parameters['policyid']

    try:
        policy = Policy.objects.get(policyid=policyid)
        Entity.objects.create(entity=entity, policyid=policy)
    except Exception as e:
        print(e)
        return Response('Cannot create the entity record', status=status.HTTP_400_BAD_REQUEST)

    return Response(status=status.HTTP_201_CREATED)

'''
    Receive a signed receipt
'''
@csrf_exempt
@api_view(('POST',))
def signReceipt(request):
    parameters = json.loads(request.body)
    print(parameters)
    idreceipt = parameters['id']
    receipt = parameters['receipt']

    try:
        stayid = Receipt.objects.get(id=idreceipt).stayid

        url = settings.RECEIPTSTORE
        r = {
            'json_receipt':receipt,
            'email': receipt['userid'],
            }
        x = requests.post(url, json=r)

        if int(x.status_code/100) != 2:
            raise Exception('RM failed')

        result = json.loads(x.text)
        print(result)

        # Store stay in data manager
        url = settings.DATA_RETENTION_STAY
        r = {
            'datein': stayid.datein.strftime('%Y-%m-%d %H:%M:%S'),
            'dateout': stayid.dateout.strftime('%Y-%m-%d %H:%M:%S'),
            'email': stayid.email.email,
            'receipt_id': result['id_receipt'],
            'data_conn': settings.DATA_CONN
            }
        x = requests.post(url, json=r)

        if int(x.status_code/100) != 2:
            raise Exception('RM failed')

        # Delete temporary receipt
        Receipt.objects.filter(id=idreceipt).delete()

        # Store receipt reference
        Stay_Receipt.objects.create(stayid=stayid, receiptid=result['id_receipt'])

    except Exception as e:
        print(e)
        return Response('ERROR', status=status.HTTP_400_BAD_REQUEST)
    return Response(status=status.HTTP_201_CREATED)

@csrf_exempt
@api_view(('DELETE',))
def requestDeletion(request):
    try:
        receiptid = request.GET['receiptid']

        # get the stay
        qs = Stay_Receipt.objects.filter(receiptid=receiptid)
        if qs.exists():
            stay = qs.first()
            stay.stayid.state = State.re
            stay.stayid.save()
        else:
            raise Exception('Stay does not exist')

    except Exception as e:
        print(e)
        return Response(f'Exception: {e}\n', status=status.HTTP_400_BAD_REQUEST)
    return Response('Deletion Requested', status=status.HTTP_200_OK)

@csrf_exempt
@api_view(('DELETE',))
def confirmDeletion(request):
    print('F()')
    try:
        stayid = request.GET['stayid']
        anonymize = bool(request.GET['anonymize'])
        print(f'{stayid} {anonymize}')

        # get receipt_id
        qs = Stay_Receipt.objects.filter(stayid=stayid)
        if qs.exists():
            stay = qs.first()
            rid = stay.receiptid

            # Confirm Data Deletion
            url = settings.DATA_DELETE
            r = {'receiptid': rid, 'anonymize':anonymize}
            x = requests.delete(url, params=r)

            if int(x.status_code/100) != 2:
                raise Exception('DM failed')
            else:
                #TODO: store summary if necessary
                print(x.text)
                stay.stayid.state = State.de
                stay.stayid.save()
        else:
            raise Exception('Stay does not exist')

    except Exception as e:
        print(e)
        return Response(f'Exception: {e}\n', status=status.HTTP_400_BAD_REQUEST)
    return Response('Deletion Requested', status=status.HTTP_200_OK)


from functools import reduce

consonants = ['b', 'c', 'd', 'f', 'g', 'h', 'j', 'k', 'l', 'm', 'n', 'p', 'q', 'r', 's', 't', 'v', 'w', 'x', 'z']
vowels = ['a', 'e', 'i', 'o', 'u', 'y']

# **Rule 4**: If a word contains a "y" after a consonant cluster or as the second letter in a two letter word it makes a vowel sound (e.g. "rhythm" -> "ythmrhay", "my" -> "ymay").
# otherwise it gets processed incorrectly.
#
# This mandates excluding 'y' from this list.
vowel_sounds = ['a', 'e', 'i', 'o', 'u', 'xr', 'yt']

def translate(text):
    """Translate the given text string to Pig Latin.
    Parameters:
    text: A string with the text to translate. Should be lower case and
    any words in a phrase separated by whitespace.
    Returns:
    Pig Latin version of the text input.
    Problems:
    * Does not handle upper case words, punctuation, whitespace between
    words is not preserved.
    """
    # List comprehension splits the word on space boundaries and calls
    # translate_word on each section. The resulting list is then glued
    # back into a string with a space between each word.
    return ' '.join([translate_word(word) for word in text.split()])

def translate_word(text):
    """Translate a single word into Pig-Latin
    Parameters:
    * text: The word to translate
    Returns
    * translated version of the input word.
    """
    # **Rule 1**: If a word begins with a vowel sound, add an "ay" sound to the end of the word. Please note that "xr" and "yt" at the beginning of a word make vowel sounds (e.g. "xray" -> "xrayay", "yttria" -> "yttriaay").
    # reduce call checks if the text starts with anything in vowel_sounds. If anything is true it will
    # return true due to or-ing everything together.
    if reduce(lambda x, y: x or text.lower().startswith(y), vowel_sounds, False):
        result = text + 'ay'
    else:
        # **Rule 2**: If a word begins with a consonant sound, move it to the end of the word and then add an "ay" sound to the end of the word. Consonant sounds can be made up of multiple consonants, a.k.a. a consonant cluster (e.g. "chair" -> "airchay").
        index = 1
        while index < len(text) and (text[index] in consonants or (index == 0 and text[index] == 'y')):
            index += 1

        # **Rule 3**: If a word starts with a consonant sound followed by "qu", move it to the end of the word, and then add an "ay" sound to the end of the word (e.g. "square" -> "aresquay").
        if index > 0 and index < len(text) and text[index] == 'u' and text[index - 1] == 'q' :
            index += 1

        result = text[index:] + text[0:index] + 'ay'
    return result


from collections import deque
import urllib2
import httplib
import operator
import socket
import threading
import urllib
import sys
import time
import ssl
import copy

import w3af.core.controllers.output_manager as om
import w3af.core.data.kb.config as cf

from w3af.core.data.constants.response_codes import NO_CONTENT
from w3af.core.controllers.exceptions import (BaseFrameworkException,
                                              ScanMustStopByKnownReasonExc)

HANDLE_ERRORS = 1 if sys.version_info < (2, 4) else 0
DEBUG = False

# Max connections allowed per host.
MAXCONNECTIONS = 50
# If found this number of timeouts in-a-row per target host then we may either:
#    1) Shrink the pool size (we might be stressing the remote host) and retry;
#        or:
#    2) Give up and make w3af stop sending requests for that host.
# This number should not be greater than (MAXCONNECTIONS / 2)
IN_A_ROW_TIMEOUTS = 15
# Constants for responses statuses
RESP_OK = 0
RESP_TIMEOUT = 1
RESP_BAD = 2

class URLTimeoutError(urllib2.URLError):
    """
    Our own URLError timeout exception. Basically a wrapper for socket.timeout.
    """
    def __init__(self):
        urllib2.URLError.__init__(self, (408, 'timeout'))

    def __str__(self):
        default_timeout = socket.getdefaulttimeout()
        if default_timeout is not None:
            return 'HTTP timeout error after %s seconds.' % default_timeout
        else:
            return 'HTTP timeout error.'

def closeonerror(read_meth):
    """
    Decorator function. When calling decorated `read_meth` if an error occurs
    we'll proceed to invoke `inst`'s close() method.
    """
    def new_read_meth(inst):
        try:
            return read_meth(inst)
        except httplib.HTTPException:
            inst.close()
            raise
    return new_read_meth

class HTTPResponse(httplib.HTTPResponse):
    # we need to subclass HTTPResponse in order to
    # 1) add readline() and readlines() methods
    # 2) add close_connection() methods
    # 3) add info() and geturl() methods

    # in order to add readline(), read must be modified to deal with a
    # buffer.  example: readline must read a buffer and then spit back
    # one line at a time.  The only real alternative is to read one
    # BYTE at a time (ick).  Once something has been read, it can't be
    # put back (ok, maybe it can, but that's even uglier than this),
    # so if you THEN do a normal read, you must first take stuff from
    # the buffer.

    # the read method wraps the original to accomodate buffering,
    # although read() never adds to the buffer.
    # Both readline and readlines have been stolen with almost no
    # modification from socket.py

    def __init__(self, sock, debuglevel=0, strict=0, method=None):
        httplib.HTTPResponse.__init__(self, sock, debuglevel, strict=strict,
                                      method=method)
        self.fileno = sock.fileno
        self.code = None
        self._rbuf = ''
        self._rbufsize = 8096
        self._handler = None  # inserted by the handler later
        self._host = None   # (same)
        self._url = None     # (same)
        self._connection = None  # (same)
        self._method = method
        self._multiread = None
        self._encoding = None

    def geturl(self):
        return self._url

    URL = property(geturl)

    def get_encoding(self):
        return self._encoding

    def set_encoding(self, enc):
        self._encoding = enc

    encoding = property(get_encoding, set_encoding)

    def _raw_read(self, amt=None):
        """
        This is the original read function from httplib with a minor
        modification that allows me to check the size of the file being
        fetched, and throw an exception in case it is too big.
        """
        if self.fp is None:
            return ''

        max_file_size = cf.cf.get('max_file_size') or None
        if max_file_size:
            if self.length > max_file_size:
                self.status = NO_CONTENT
                self.reason = 'No Content'  # Reason-Phrase
                self.close()
                return ''

        if self.chunked:
            return self._read_chunked(amt)

        if amt is None:
            # unbounded read
            if self.length is None:
                s = self.fp.read()
            else:
                s = self._safe_read(self.length)
                self.length = 0
            self.close()        # we read everything
            return s

        if self.length is not None:
            if amt > self.length:
                # clip the read to the "end of response"
                amt = self.length

        # we do not use _safe_read() here because this may be a .will_close
        # connection, and the user is reading more bytes than will be provided
        # (for example, reading in 1k chunks)
        s = self.fp.read(amt)
        if self.length is not None:
            self.length -= len(s)

        return s

    def close(self):
        # First call parent's close()
        httplib.HTTPResponse.close(self)
        if self._handler:
            self._handler._request_closed(self._connection)

    def close_connection(self):
        self._handler._remove_connection(self._host, self._connection)
        self.close()

    def info(self):
        return self.headers

    @closeonerror
    def read(self, amt=None):
        # w3af does always read all the content of the response...
        # and I also need to do multiple reads to this response...
        # BUGBUG: Is this OK? What if a HEAD method actually returns something?!
        if self._method == 'HEAD':
            # This indicates that we have read all that we needed from the socket
            # and that the socket can be reused!
            #
            # This like fixes the bug with title "GET is much faster than HEAD".
            #https://sourceforge.net/tracker2/?func=detail&aid=2202532&group_id=170274&atid=853652
            self.close()
            return ''

        if self._multiread is None:
            #read all
            self._multiread = self._raw_read()

        if not amt is None:
            L = len(self._rbuf)
            if amt > L:
                amt -= L
            else:
                s = self._rbuf[:amt]
                self._rbuf = self._rbuf[amt:]
                return s
        else:
            s = self._rbuf + self._multiread
            self._rbuf = ''
            return s

    def readline(self, limit=-1):
        i = self._rbuf.find('\n')
        while i < 0 and not (0 < limit <= len(self._rbuf)):
            new = self._raw_read(self._rbufsize)
            if not new:
                break
            i = new.find('\n')
            if i >= 0:
                i = i + len(self._rbuf)
            self._rbuf = self._rbuf + new
        if i < 0:
            i = len(self._rbuf)
        else:
            i = i + 1
        if 0 <= limit < len(self._rbuf):
            i = limit
        data, self._rbuf = self._rbuf[:i], self._rbuf[i:]
        return data

    @closeonerror
    def readlines(self, sizehint=0):
        total = 0
        line_list = []
        while 1:
            line = self.readline()
            if not line:
                break
            line_list.append(line)
            total += len(line)
            if sizehint and total >= sizehint:
                break
        return line_list

    def set_body(self, data):
        """
        This was added to make my life a lot simpler while implementing mangle
        plugins
        """
        self._multiread = data

class ConnectionManager(object):
    """
    The connection manager must be able to:
        * keep track of all existing HTTPConnections
        * kill the connections that we're not going to use anymore
        * Create/reuse connections when needed.
        * Control the size of the pool.
    """

    def __init__(self):
        self._lock = threading.RLock()
        self._host_pool_size = MAXCONNECTIONS
        self._hostmap = {}  # map hosts to a list of connections
        self._used_cons = []  # connections being used per host
        self._free_conns = []  # available connections

    def remove_connection(self, conn, host=None):
        """
        Remove a connection, it was closed by the server.

        :param conn: Connection to remove
        :param host: The host for to the connection. If passed, the connection
        will be removed faster.
        """
        # Just make sure we don't leak open connections
        conn.close()

        with self._lock:

            if host:
                if host not in self._hostmap:
                    raise ValueError('Host "%s" not present in pool.' % host)
                if conn in self._hostmap[host]:
                    self._hostmap[host].remove(conn)

            else:
                # We don't know the host. Need to find it by looping
                for _host, conns in self._hostmap.items():
                    if conn in conns:
                        host = _host
                        conns.remove(conn)
                        break

            for lst in (self._free_conns, self._used_cons):
                try:
                    lst.remove(conn)
                except ValueError:
                    # I don't care much about the connection not being in
                    # the free_conns or used_conns. This might happen because
                    # of a thread locking issue (basically, someone is not
                    # locking before moving connections around).
                    pass

            # No more conns for 'host', remove it from mapping
            conn_total = self.get_connections_total(host)
            if host and not conn_total:
                del self._hostmap[host]

            if DEBUG:
                msg = 'keepalive: removed one connection, len(self._hostmap' \
                      '["%s"]): %s' % (host, conn_total)
                om.out.debug(msg)

    def free_connection(self, conn):
        """
        Recycle a connection. Mark it as available for being reused.
        """
        if conn in self._used_cons:
            self._used_cons.remove(conn)
            self._free_conns.append(conn)

    def replace_connection(self, bad_conn, host, conn_factory):
        """
        Re-create a mal-functioning connection.

        :param bad_conn: The bad connection
        :param host: The host for the connection
        :param conn_factory: The factory function for new connection creation.
        """
        with self._lock:
            self.remove_connection(bad_conn, host)
            if DEBUG:
                msg = 'keepalive: replacing bad connection with a new one'
                om.out.debug(msg)
            new_conn = conn_factory(host)
            conns = self._hostmap.setdefault(host, [])
            conns.append(new_conn)
            self._used_cons.append(new_conn)
            return new_conn

    def get_available_connection(self, host, conn_factory):
        """
        Return an available connection ready to be reused

        :param host: Host for the connection.
        :param conn_factory: Factory function for connection creation. Receives
            <host> as parameter.
        """
        with self._lock:
            retry_count = 10

            while retry_count > 0:
                ret_conn = None

                # First check if we can reuse an existing free conn.
                for conn in self._hostmap.setdefault(host, []):
                    try:
                        self._free_conns.remove(conn)
                    except ValueError:
                        continue
                    else:
                        self._used_cons.append(conn)
                        ret_conn = conn
                        break

                # No?... well, let's try to create a new one.
                conn_total = self.get_connections_total(host)
                if not ret_conn and conn_total < self._host_pool_size:
                    if DEBUG:
                        msg = 'keepalive: added one connection, len(self._hostmap'\
                              '["%s"]): %s' % (host, conn_total + 1)
                        om.out.debug(msg)
                    ret_conn = conn_factory(host)
                    self._used_cons.append(ret_conn)
                    self._hostmap[host].append(ret_conn)

                if ret_conn is not None:  # Good! We have one!
                    return ret_conn
                else:  # Maybe we should wait a little and try again 8^)
                    retry_count -= 1
                    time.sleep(0.3)

            msg = 'keepalive: been waiting too long for a pool connection.' \
                  ' I\'m giving up. Seems like the pool is full.'

            if DEBUG:
                om.out.debug(msg)

            raise BaseFrameworkException(msg)

    def resize_pool(self, new_size):
        """
        Set a new pool size.
        """
        pass

    def get_all(self, host=None):
        """
        If <host> is passed return a list containing the created connections
        for that host. Otherwise return a dict with 'host: str' and
        'conns: list' as items.

        :param host: Host
        """
        if host:
            return list(self._hostmap.get(host, []))
        else:
            return dict(self._hostmap)

    def get_connections_total(self, host=None):
        """
        If <host> is None return the grand total of created connections;
        otherwise return the total of created conns. for <host>.
        """
        values = self._hostmap.values() if (host is None) \
            else [self._hostmap[host]]
        return reduce(operator.add, map(len, values or [[]]))

class KeepAliveHandler(object):

    def __init__(self):
        # Create the connection pool instance
        #
        # Note: In the initial code this connection manager was created at
        #       the module level and was shared between the HTTP and HTTPS
        #       keep alive handlers. This was buggy since at any point if
        #       a user requested http://host.tld and then https://host.tld
        #       a connection to the HTTP one was returned from the manager
        #       for the HTTPS request.
        #
        #       This change lets us still keep a "persistent" connection
        #       manager since our opener_settings and xurllib will only
        #       create one instance for the KeepAliveHandler and use that
        #       during the whole scan.
        self._cm = ConnectionManager()

        # Typically a urllib2.OpenerDirector instance. Set by the
        # urllib2 mechanism.
        self.parent = None
        self._pool_lock = threading.RLock()
        # Map hosts to a `collections.deque` of response status.
        self._hostresp = {}
        # Current number of 'in-a-row-timeouts'. It may vary depending on the
        # current size of the pool.
        self._curr_check_failures = IN_A_ROW_TIMEOUTS
        # Tail list filter factory function
        self._get_tail_filter = lambda: deque(maxlen=self._curr_check_failures)

    def get_open_connections(self):
        """
        Return a list of connected hosts and the number of connections
        to each.  [('foo.com:80', 2), ('bar.org', 1)]
        """
        return [(host, len(li)) for (host, li) in self._cm.get_all().items()]

    def close_connection(self, host):
        """
        Close connection(s) to <host>
        host is the host:port spec, as in 'www.cnn.com:8080' as passed in.
        no error occurs if there is no connection to that host.
        """
        for conn in self._cm.get_all(host):
            self._cm.remove_connection(conn, host)

    def close_all(self):
        """
        Close all open connections
        """
        for conns in self._cm.get_all().values():
            for conn in conns:
                self._cm.remove_connection(conn)

    def _request_closed(self, connection):
        """
        Tells us that this request is now closed and that the
        connection is ready for another request
        """
        self._cm.free_connection(connection)

    def _remove_connection(self, host, conn):
        self._cm.remove_connection(conn, host)

    def do_open(self, req):
        """
        Called by handler's url_open method.
        """
        host = req.get_host()
        if not host:
            raise urllib2.URLError('no host given')

        try:
            resp_statuses = self._hostresp.setdefault(host,
                                                      self._get_tail_filter())
            # Check if all our last 'resp_statuses' were timeouts and raise
            # a ScanMustStopException if this is the case.
            if len(resp_statuses) == self._curr_check_failures:

                # https://mail.python.org/pipermail/python-dev/2007-January/070515.html
                # https://github.com/andresriancho/w3af/search?q=deque&ref=cmdform&type=Issues
                # https://github.com/andresriancho/w3af/issues/1311
                #
                # I copy the deque to avoid issues when iterating over it in the
                # all() / for statement below
                resp_statuses_cp = copy.copy(resp_statuses)

                if all(st == RESP_TIMEOUT for st in resp_statuses_cp):
                    msg = ('w3af found too many consecutive timeouts. The'
                           ' remote webserver seems to be unresponsive; please'
                           ' verify manually.')
                    reason = 'Timeout while trying to reach target.'
                    raise ScanMustStopByKnownReasonExc(msg, reason=reason)

            conn_factory = self._get_connection
            conn = self._cm.get_available_connection(host, conn_factory)

            if conn.is_fresh:
                # First of all, call the request method. This is needed for
                # HTTPS Proxy
                if isinstance(conn, ProxyHTTPConnection):
                    conn.proxy_setup(req.get_full_url())

                conn.is_fresh = False
                self._start_transaction(conn, req)
                resp = conn.getresponse()
            else:
                # We'll try to use a previously created connection
                resp = self._reuse_connection(conn, req, host)
                # If the resp is None it means that connection is bad. It was
                # possibly closed by the server. Replace it with a new one.
                if resp is None:
                    conn.close()
                    conn = self._cm.replace_connection(conn, host,
                                                       conn_factory)
                    # First of all, call the request method. This is needed for
                    # HTTPS Proxy
                    if isinstance(conn, ProxyHTTPConnection):
                        conn.proxy_setup(req.get_full_url())

                    # Try again with the fresh one
                    conn.is_fresh = False
                    self._start_transaction(conn, req)
                    resp = conn.getresponse()

        except socket.timeout:
            # We better discard this connection
            self._cm.remove_connection(conn, host)
            resp_statuses.append(RESP_TIMEOUT)
            raise URLTimeoutError()

        except socket.error:
            # We better discard this connection
            self._cm.remove_connection(conn, host)
            resp_statuses.append(RESP_BAD)
            raise

        except httplib.HTTPException:
            # We better discard this connection
            self._cm.remove_connection(conn, host)
            resp_statuses.append(RESP_BAD)
            raise

        else:
            # This response seems to be fine
            resp_statuses.append(RESP_OK)

            # If not a persistent connection, don't try to reuse it
            if resp.will_close:
                self._cm.remove_connection(conn, host)

            if DEBUG:
                om.out.debug("STATUS: %s, %s" % (resp.status, resp.reason))

            resp._handler = self
            resp._host = host
            resp._url = req.get_full_url()
            resp._connection = conn
            resp.code = resp.status
            resp.headers = resp.msg
            resp.msg = resp.reason
            return resp

    def _reuse_connection(self, conn, req, host):
        """
        Start the transaction with a re-used connection
        return a response object (r) upon success or None on failure.
        This DOES not close or remove bad connections in cases where
        it returns.  However, if an unexpected exception occurs, it
        will close and remove the connection before re-raising.
        """
        try:
            self._start_transaction(conn, req)
            r = conn.getresponse()
            # note: just because we got something back doesn't mean it
            # worked.  We'll check the version below, too.
        except (socket.error, httplib.HTTPException):
            r = None
        except:
            # adding this block just in case we've missed something we will
            # still raise the exception, but lets try and close the connection
            # and remove it first.  We previously got into a nasty loop where
            # an exception was uncaught, and so the connection stayed open.
            # On the next try, the same exception was raised, etc. The tradeoff
            # is that it's now possible this call will raise a DIFFERENT
            # exception
            if DEBUG:
                om.out.error("unexpected exception - closing connection to %s"
                             " (%d)" % (host, id(conn)))
            self._cm.remove_connection(conn, host)
            raise

        if r is None or r.version == 9:
            # httplib falls back to assuming HTTP 0.9 if it gets a
            # bad header back.  This is most likely to happen if
            # the socket has been closed by the server since we
            # last used the connection.
            if DEBUG:
                om.out.debug("failed to re-use connection to %s (%d)" % (host,
                             id(conn)))
            r = None
        else:
            if DEBUG:
                om.out.debug(
                    "re-using connection to %s (%d)" % (host, id(conn)))
            r._multiread = None

        return r

    def _start_transaction(self, conn, req):
        """
        The real workhorse.
        """
        try:
            data = req.get_data()
            if data is not None:
                data = str(data)
                conn.putrequest(req.get_method(), req.get_selector(),
                                skip_host=1, skip_accept_encoding=1)

                if not req.has_header('Content-type'):
                    conn.putheader('Content-type',
                                   'application/x-www-form-urlencoded')

                if not req.has_header('Content-length'):
                    conn.putheader('Content-length', '%d' % len(data))
            else:
                conn.putrequest(req.get_method(), req.get_selector(),
                                skip_host=1, skip_accept_encoding=1)
        except (socket.error, httplib.HTTPException):
            raise
        else:
            # Add headers
            header_dict = dict(self.parent.addheaders)
            header_dict.update(req.headers)
            header_dict.update(req.unredirected_hdrs)

            for k, v in header_dict.iteritems():
                conn.putheader(to_utf8_raw(k),
                               to_utf8_raw(v))
            conn.endheaders()

            if data is not None:
                conn.send(data)

    def _get_connection(self, host):
        """
        "Abstract" method.
        """
        raise NotImplementedError()

class HTTPHandler(KeepAliveHandler, urllib2.HTTPHandler):
    def __init__(self):
        KeepAliveHandler.__init__(self)

    def http_open(self, req):
        return self.do_open(req)

    def _get_connection(self, host):
        return HTTPConnection(host)

class HTTPSHandler(KeepAliveHandler, urllib2.HTTPSHandler):
    def __init__(self, proxy):
        KeepAliveHandler.__init__(self)
        self._proxy = proxy
        host = port = ''
        try:
            host, port = self._proxy.split(':')
        except:
            msg = 'The proxy you are specifying (%s) is invalid! The expected'\
                  ' format is <ip_address>:<port> is expected.'
            raise BaseFrameworkException(msg % proxy)

        if not host or not port:
            self._proxy = None

    def https_open(self, req):
        return self.do_open(req)

    def _get_connection(self, host):
        if self._proxy:
            proxy_host, port = self._proxy.split(':')
            return ProxyHTTPSConnection(proxy_host, port)
        else:
            return HTTPSConnection(host)

class _HTTPConnection(httplib.HTTPConnection):

    def __init__(self, host, port=None, strict=None,
                 timeout=socket._GLOBAL_DEFAULT_TIMEOUT):
        httplib.HTTPConnection.__init__(self, host, port, strict,
                                        timeout=timeout)
        self.is_fresh = True

class ProxyHTTPConnection(_HTTPConnection):
    """
    this class is used to provide HTTPS CONNECT support.
    """
    _ports = {'http': 80, 'https': 443}

    def __init__(self, host, port=None, strict=None):
        _HTTPConnection.__init__(self, host, port, strict)

    def proxy_setup(self, url):
        #request is called before connect, so can interpret url and get
        #real host/port to be used to make CONNECT request to proxy
        proto, rest = urllib.splittype(url)
        if proto is None:
            raise ValueError("unknown URL type: %s" % url)
        #get host
        host, rest = urllib.splithost(rest)
        self._real_host = host

        #try to get port
        host, port = urllib.splitport(host)
        #if port is not defined try to get from proto
        if port is None:
            try:
                self._real_port = self._ports[proto]
            except KeyError:
                raise ValueError("unknown protocol for: %s" % url)
        else:
            self._real_port = int(port)

    def connect(self):
        httplib.HTTPConnection.connect(self)

        #send proxy CONNECT request
        new_line = '\r\n'
        self.send("CONNECT %s:%d HTTP/1.1%s" % (self._real_host,
                                                self._real_port,
                                                new_line))

        connect_headers = {'Proxy-Connection': 'keep-alive',
                           'Connection': 'keep-alive',
                           'Host': self._real_host}

        for header_name, header_value in connect_headers.items():
            self.send('%s: %s%s' % (header_name, header_value, new_line))

        self.send(new_line)

        #expect a HTTP/1.0 200 Connection established
        response = self.response_class(self.sock, strict=self.strict,
                                       method=self._method)
        (version, code, message) = response._read_status()
        #probably here we can handle auth requests...
        if code != 200:
            #proxy returned and error, abort connection, and raise exception
            self.close()
            raise socket.error("Proxy connection failed: %d %s" %
                              (code, message.strip()))
        #eat up header block from proxy....
        while True:
            #should not use directly fp probably
            line = response.fp.readline()
            if line == '\r\n':
                break

class ProxyHTTPSConnection(ProxyHTTPConnection):
    """
    this class is used to provide HTTPS CONNECT support.
    """
    default_port = 443

    # Customized response class
    response_class = HTTPResponse

    def __init__(self, host, port=None, key_file=None, cert_file=None,
                 strict=None):
        ProxyHTTPConnection.__init__(self, host, port)
        self.key_file = key_file
        self.cert_file = cert_file

    def connect(self):
        ProxyHTTPConnection.connect(self)
        #make the sock ssl-aware
        ssl_sock_inst = ssl.wrap_socket(self.sock, self.key_file,
                                        self.cert_file)
        self.sock = ssl_sock_inst

def to_utf8_raw(unicode_or_str):
    if isinstance(unicode_or_str, unicode):
        # Is 'ignore' the best option here?
        return unicode_or_str.encode('utf-8', 'ignore')
    return unicode_or_str

class HTTPConnection(_HTTPConnection):
    # use the modified response class
    response_class = HTTPResponse

    def __init__(self, host, port=None, strict=None):
        _HTTPConnection.__init__(self, host, port, strict)

class HTTPSConnection(httplib.HTTPSConnection):
    response_class = HTTPResponse

    def __init__(self, host, port=None, key_file=None, cert_file=None,
                 strict=None):
        httplib.HTTPSConnection.__init__(self, host, port, key_file, cert_file,
                                         strict)
        self.is_fresh = True


from faker import Faker
from freezegun import freeze_time

from doajtest.fixtures import AccountFixtureFactory, JournalFixtureFactory, BackgroundFixtureFactory
from doajtest.helpers import DoajTestCase
from portality import models
from portality.core import app
from portality.lib import anon
from portality.lib.anon import basic_hash, anon_name, anon_email
from portality.tasks import anon_export

class TestAnon(DoajTestCase):
    def setUp(self):
        self.old_anon_salt = app.config['ANON_SALT']
        app.config['ANON_SALT'] = 'testsalt'
        anon.fake = Faker()
        anon.fake.seed(1)

    def tearDown(self):
        app.config['ANON_SALT'] = self.old_anon_salt

    def test_01_basic_hash(self):
        assert basic_hash('test content') == '259ea4fab4e03a26c25f2b55f37a2f571931797d67f033e8898e76481d2a8563', basic_hash('test content')

    def test_02_anon_name(self):
        assert anon_name() == 'Ryan Gallagher', anon_name()

    def test_03_anon_email(self):
        assert anon_email('test@doaj.org') == '508dd70a7c888d9985c5ed37276d1138d73db171932b5866d48f581dc6119ac5@example.com'

    def test_04_anonymise_email(self):
        record = models.Account(**AccountFixtureFactory.make_publisher_source())
        e = anon_export._anonymise_email(record).email
        assert e == '25011d8de5bfcb72ee529fcc38b518ea6a46f99a81a412c065fe7147272b8f2a@example.com', e

    def test_05_anonymise_admin_with_notes(self):
        journal_src = JournalFixtureFactory.make_journal_source()

        journal_src['admin'] = {
            'owner': 'testuser',
            'editor': 'testeditor',
            'notes': [
                {
                    "id" : "note1",
                    'note': 'Test note',
                    'date': '2017-02-23T00:00:00Z'
                },
                {
                    "id" : "note2",
                    'note': 'Test note 2',
                    'date': '2017-02-23T00:00:00Z'
                }
            ]
        }

        journal = models.Journal(**journal_src)

        with freeze_time("2017-02-23"):
            ar = anon_export._anonymise_admin(journal)

        assert ar.data['admin'] == {
            'owner': 'testuser',
            'editor': 'testeditor',
            'notes': [
                {
                    "id" : "note1",
                    'note': 'f4007b0953d4a9ecb7e31820b5d481d96ee5d74a0a059a54f07a326d357ed895',
                    'date': '2017-02-23T00:00:00Z'
                },
                {
                    "id" : "note2",
                    'note': '772cf6f91219db969e4aa28e4fd606b92316948545ad528fd34feb1b9b12a3ad',
                    'date': '2017-02-23T00:00:00Z'
                }
            ]
        }, ar['admin']

    def test_06_anonymise_admin_empty_notes(self):
        journal_src = JournalFixtureFactory.make_journal_source()

        journal_src['admin'] = {
            'owner': 'testuser',
            'editor': 'testeditor',
            'notes': []
        }

        journal = models.Journal(**journal_src)

        with freeze_time("2017-02-23"):
            ar = anon_export._anonymise_admin(journal)

        assert ar.data['admin'] == {
            'owner': 'testuser',
            'editor': 'testeditor',
            'notes': []
        }, ar['admin']

    def test_07_anonymise_account(self):
        anon_a = anon_export.anonymise_account(AccountFixtureFactory.make_publisher_source())
        assert anon_a['id'] == 'publisher', anon_a['id']
        assert anon_a['email'] == '25011d8de5bfcb72ee529fcc38b518ea6a46f99a81a412c065fe7147272b8f2a@example.com', anon_a['email']

    def test_10_anonymise_background_job(self):
        bgjob = BackgroundFixtureFactory.example()
        bgjob['params'].update({'suggestion_bulk_edit__note': 'Test note'})
        assert anon_export.anonymise_background_job(bgjob)['params'] == {'suggestion_bulk_edit__note': 'f4007b0953d4a9ecb7e31820b5d481d96ee5d74a0a059a54f07a326d357ed895'}


from __future__ import print_function
"""
JS9 wrapper to be used in Jupyter/IPython notebooks

Matias Carrasco Kind
"""
__author__ = 'Matias Carrasco Kind'
import uuid
import json
from astropy import wcs
from astropy.io import fits
import pyjs9
from six import BytesIO
import base64
from socketIO_client import SocketIO
import re
import os

init_js9 = False

def libinit(root='http://localhost:8888/files/'):
    """
    Initialization function needed to load javascript libraries and css styles into the notebook.

    Parameters:
    -----------
    root : root url where main path for the notebook is running
    """
    global init_js9
    init_js9 = True
    #root = 'http://localhost:8888/files/'
    req = """
    require.config({
      paths: {
          js9: '%sjs9-allinone',
      }
    });
    """ % root
    css = """
    <link rel="stylesheet" type="text/css" href="%sjs9-allinone.css">
    """ % root
    init = """
    require(['js9'], function(){
        console.log('Loaded');
    });
    """
    get_ipython().run_cell_magic('javascript', '', req)
    get_ipython().run_cell_magic('html', '', css)
    get_ipython().run_cell_magic('javascript', '', init)

class Js9Local(object):
    """Jupyter Notebbok wrapper around JS9 class object"""
    def __init__(self, root='http://localhost:8888/files/',do_init=False):
        """
        Initiate the class with an unique id unless is manually parsed

        Parameters:
        -----------
        wid : Id to be use for the JS9 display. Default a unique id is created
        """
        global init_js9
        if not init_js9:
             if do_init:libinit(root=root)
             init_js9 = True

        self.created = False
        self.wid = uuid.uuid4().hex
        self.url = None
        self.wcs = None
        self.header = None
        self.suffix = 'JS9'

    def DelDiv(self):
        """
        Delete current div for this class
        """
        if self.created:
            self.CloseImage()
            command = """$('#{}').remove();""".format(self.wid)
            get_ipython().run_cell_magic('javascript', '', command)
            self.created = False
            self.wid = uuid.uuid4().hex

    def NewDiv(self, width='100%', height='512px'):
        """
        Creates a new div to be added to the notebook cell

        Parameters:
        -----------

        width :  width of the displayed window in css style (string). Default is
                 80% of the screen
        height : height of the displayed window in css style (string). Default is
                 512px of the screen

        """
        if self.created:
            self.DelDiv()
        print('Display id = {}JS9'.format(self.wid))
        command = """
        "

"
        """.format(self.wid, width, height)
        js_command = "element.append({});".format(command)
        get_ipython().run_cell_magic('javascript', '', js_command)
        get_ipython().run_cell_magic('javascript', '', "JS9.AddDivs('{}JS9');".format(self.wid))
        self.created = True

    def Load(self, url, **kwargs):
        """
        Loads a FITS file into the display created.
        Parameters:
        -----------
        url : Path to file (can be locally)
        Extra arguments in form of a dictionary used by the JS9 javascript object JS9.Load()
        """
        opts = ''
        self.url = url
        if not 'wcsunits' in kwargs:
            kwargs['wcsunits'] = 'degrees'
        if not 'scaleclipping' in kwargs:
            kwargs['scaleclipping'] = 'zscale'
        if len(kwargs) > 0:
            opts = json.dumps(kwargs)+','
        fmt = dict(url=self.url, kw=opts, wid=self.wid, suffix=self.suffix)
        command = "JS9.Load('{url}',{kw}{{display:'{wid}{suffix}'}});".format(**fmt)
        try:
            hdulist = fits.open(url)
            temp_wcs = wcs.WCS(hdulist[0].header)
            self.header = hdulist[0].header
            self.wcs_world2pix = temp_wcs.wcs_world2pix
            self.wcs_pix2world = temp_wcs.wcs_pix2world
            hdulist.close()
        except:
            pass
        get_ipython().run_cell_magic('javascript', '', command)

    def CloseImage(self):
        """
        Close current image
        """
        command = "JS9.CloseImage({{display:'{wid}{suffix}'}});".format(wid=self.wid, suffix=self.suffix)
        get_ipython().run_cell_magic('javascript', '', command)

    def SetColorMap(self, colormap, contrast=None, bias=None):
        """
        Set color map of currently displayed image

        Parameters:
        -----------

        colormap: colormap name
        contrast: contrast value (range: 0 to 10)
        bias: bias value (range 0 to 1)
        """
        extra = ''
        if contrast is not None:
            extra += '%f,' % contrast
        if bias is not None:
            extra  += '%f,' % bias
        fmt = dict(wid=self.wid,cmap=colormap, extra=extra, suffix=self.suffix)
        command = "JS9.SetColormap('{cmap}', {extra} {{display:'{wid}{suffix}'}});".format(**fmt)
        get_ipython().run_cell_magic('javascript', '', command)

    def AddRegions(self, **kwargs):
        """
        Add Regions to JS9 display using same syntax as JS9.AddRegions.
        This uses astropy.wcs to do the coordinates transformations
        """
        # Addregions use pixel coordinates. listRegions and SaveRegions use RA and Dec.
        n_objs = 0
        objs = []
        # default shape is circle
        if not 'shape' in kwargs:
            kwargs['shape'] = ['circle']
        for k in kwargs.keys():
            n_objs = max(n_objs, len(kwargs[k]))
        for j in range(n_objs):
            temp = {}
            for k in kwargs.keys():
                try:
                    temp[k] = kwargs[k][j]
                except IndexError:
                    if k == 'shape':
                        temp[k] = 'circle'
            objs.append(temp)
        self.all_objs = json.dumps(objs)
        command = "JS9.AddRegions({objs}, {{display:'{wid}{suffix}'}})".format(objs=self.all_objs, wid=self.wid, suffix=self.suffix)
        get_ipython().run_cell_magic('javascript', '', command)

    def RegionList(self):
        """
        Name the JS9 GetShapes object string as RegionList.
        RegionList can be called in python after running Js9Local.RegionList().
        """
        command = """
        IPython.notebook.kernel.execute("RegionList=" + JSON.stringify(JS9.GetShapes("regions", {{display: '{wid}JS9'}})));
        """.format(wid=self.wid)
        get_ipython().run_cell_magic('javascript', '', command)

    def ReadRegions(self, fname="temp"):
        """
        ReadRegions() must be be run before SaveRegions() to save regions.
        Create a default temp file with regions written in one line, which is deleted after SaveRegions().
        IPython.notebook.kernel.execute has a problem with new lines.
        For loop must be in one line for IPython.notebook.kernel.execute.
        """
        self.fname = fname
        command = """IPython.notebook.kernel.execute('file = open("temp", "w"); [file.write(x["wcsstr"]) for x in '+ JSON.stringify(JS9.GetShapes("regions", {{display: '{wid}JS9'}})) +']; file.close()');""".format(wid=self.wid)
        get_ipython().run_cell_magic('javascript', '', command)

    def SaveRegions(self, fname="js9.reg"):
        """
        Save regions to a text file in the directory of the running Jupyter notebook
        following JS9 SaveRegions format with header and each region starts with a new line.
        Default file name is js9.reg. Users can specifiy their own file name.
        """
        self.fname = fname
        temp = open('temp','r').read()
        regions = re.sub(r'(circle)', r'\n\1', temp)
        regions = "# Region file format: JS9 version 1.0\nICRS"+regions
        file = open(fname, "w")
        file.write(regions)
        file.close()
        os.remove("temp")

    def LoadRegions(self, fname='js9.reg'):
        """
        LoadRegions() loads saved regions from a user-specified file in the current Jupyter notebook's directory
        """
        self.fname=fname
        command = "JS9.LoadRegions('{fname}', {{display:'{wid}JS9'}})".format(fname=fname, wid=self.wid)
        get_ipython().run_cell_magic('javascript', '', command)

default_root = 'http://141.142.236.170'
default_port_html = 8000
default_port_io = 8001
default_width = 970
default_height= 600

def NewDiv(width=default_width, height=default_height):
    """
    Creates a new div to be added to the notebook cell
    The name 'NewDiv' is temporary.

    Parameters:
    -----------
    width :  width of the displayed window in css style (string). Default is 970.
    height : height of the displayed window in css style (string). Default is 600.
    -----------
    examples:
    >>> import jjs9
    >>> jjs9.NewDiv()
    """

    global wid
    wid = uuid.uuid4().hex
    print('Display id = {}JS9'.format(wid))
    fmt = dict(url=default_root, port0=default_port_html, wid=wid, width=width, height=height)
    html_command = """

    """.format(**fmt)
    get_ipython().run_cell_magic('html', '', html_command)

class Js9Server(pyjs9.JS9):
    """
    Connect to server that runs js9Helper.js for server-side analysis
    Js9Server is best for Jupyter notebook running on a local computer.
    Analyses are done by communicating with the JS9 back-end server using socket.io.

    Example:
    >>> J = jjs9.Js9Server()
    >>> J.LoadFITS('filename.fits')

    Note that the name assignment J = jjs9.Js9Server() will only give the name to the latest NewDiv,
    so it's recommended to give a name of a NewDiv immediately after calling NewDiv().

    All pyjs9 functionality are available. https://github.com/ericmandel/pyjs9

    Examples:
    >>> J.LoadProxy('http://hea-www.cfa.harvard.edu/~eric/coma.fits', {'scale':'linear', 'colormap':'sls'})
    >>> J.SetColormap('red')

    """

    def __init__(self, root=default_root, port_html=default_port_html, port_io=default_port_io):
        self.wcs = None
        self.header = None
        self.root = root
        self.port_html = port_html
        self.port_io = port_io
        super(Js9Server, self).__init__(host=root+':'+str(port_io),id=wid+'JS9')

    def LoadFITS(self, name=None):
        """
        Load('filename') from pyjs9 opens files on the JS9 website server.
        LoadFITS('filename.fits') opens fits files local to the user's running jupyter notebook.
        LoadFITS and SetFITS use a similar algorithm.
        """
        F=fits.open(name)
        # in-memory string
        memstr=BytesIO()
        # write fits to memory string
        F.writeto(memstr)
        # get memory string as an encoded string
        encstr = base64.b64encode(memstr.getvalue())
        # set up JS9 options
        opts = {}
        if name:
            opts['filename'] = name
        # send encoded file to JS9 for display
        got = self.Load(encstr, opts)
        # finished with memory string
        memstr.close()
        return got


import unittest

from fastapi.testclient import TestClient

from main import app

EDGE_UA = "Mozilla/5.0 (Linux; Android 10; HD1913) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/102.0.5005.58 Mobile Safari/537.36 EdgA/100.0.1185.50"

client = TestClient(app)

class BaseUnitTests(unittest.TestCase):
    def test_json(self):
        response = client.get(
            "info", params={"format": "json"}, headers={"User-Agent": EDGE_UA}
        )
        self.assertEqual(response.json(), {"user_agent": EDGE_UA})

    def test_html(self):
        response = client.get(
            "info", params={"format": "html"}, headers={"User-Agent": EDGE_UA}
        )
        self.assertIn(
            f'<input type="text" id=user-agent name=agent value="{EDGE_UA}">',
            response.content.decode(),
        )

    def test_incorrect_format(self):
        response = client.get(
            "info", params={"format": "aujfdhs"}, headers={"User-Agent": EDGE_UA}
        )
        self.assertEqual(response.status_code, 400)

    def test_no_format(self):
        response = client.get("info", headers={"User-Agent": EDGE_UA})
        self.assertEqual(response.status_code, 400)

if __name__ == "__main__":
    unittest.main()


import math
from pvti_offroad.common.utils import b_to_g_rot, wrap_to_pi_torch
from pvti_offroad.common.pytypes import VehicleState
import torch

class VehicleModel:
    def __init__(self, dt = 0.2, N_node = 10, map_info = None, n_sample = 100):
        self.local_map = map_info
        self.N = N_node
        self.m = 25
        self.width = 0.45
        self.L = 0.9
        self.Lr = 0.45
        self.Lf = 0.45
        self.Caf = self.m * self.Lf/self.L * 0.5 * 0.35 * 180/torch.pi
        self.Car = self.m * self.Lr/self.L * 0.5 * 0.35 * 180/torch.pi
        self.Izz = self.Lf*self.Lr*self.m
        self.g= 9.814195
        self.h = 0.15
        self.dt = dt
        self.delta_rate_max = 50*math.pi/180.0 # rad/s
        self.rollover_cost_scale = 10.0
        self.Q = torch.tensor([0.5, 0.5, 0.5])
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        ## dyn_state = [x,y,z,vx,vy,wz,roll,pitch,yaw]
        self.state_dim = 9
        ## inputs = [ax, delta]
        self.input_dim = 2
        self.horizon = self.dt* N_node
        # delta and ax
        self.ax_max = 1.0
        self.ax_min = -1.0
        self.delta_min = -1* 25*torch.pi/180.0
        self.delta_max = 25*torch.pi/180.0

        self.cur_state = VehicleState()

    def compute_slip(self, x,u):
        try: # when u is for batch
            delta = u[:, 1]
            clip_vx = torch.max(torch.hstack((x[:,3].view(-1,1),torch.ones(len(x[:,3])).to(device=self.device).view(-1,1))),dim =1).values
            alpha_f = delta - (x[:,4]+self.Lf*x[:,5])/clip_vx
            alpha_r = (-x[:,4]+self.Lr*x[:,5])/clip_vx
        except IndexError: # when u is single
            delta = u[1]
            vx = x[3]
            alpha_f = delta - (x[4]+self.Lf*x[5])/vx
            alpha_r = (-x[4]+self.Lr*x[5])/vx

        return alpha_f, alpha_r

    def compute_normal_force(self,x,u,roll,pitch):
        try:
            ax = u[:, 0]  # when u is for batch
        except IndexError:
            ax = u[0]  # when u is single
            ax = torch.tensor(ax)
        Fzf = self.Lr*self.m*self.g*torch.cos(pitch)*torch.cos(roll)/self.L + self.h*self.m/self.L*(ax+self.g*torch.sin(pitch))
        Fzr = self.Lr*self.m*self.g*torch.cos(pitch)*torch.cos(roll)/self.L - self.h*self.m/self.L*(ax+self.g*torch.sin(pitch))
        return Fzf, Fzr

    def dynamics_predict(self,cur_state:VehicleState, batch_u = None):
        x =  cur_state.odom.pose.pose.position.x
        y =  cur_state.odom.pose.pose.position.y
        psi = cur_state.euler.yaw
        vx = cur_state.local_twist.linear.x
        vy = cur_state.local_twist.linear.y
        wz = cur_state.local_twist.angular.z
        z =  cur_state.odom.pose.pose.position.z
        roll = cur_state.euler.roll
        pitch = cur_state.euler.pitch

        x =torch.tensor([x,y,psi,vx,vy,wz,z,roll,pitch]).to(device=self.device)

        if batch_u.shape[0] > 1:
            x = x.repeat(batch_u.shape[0],1)
        else:
            x = x.unsqueeze(dim=0)

        usets = batch_u.to(device=self.device)

        # return self.torch_dynamics_predict(cur_state = x, batch_u = usets)
        return self.torch_2d_predict(cur_state = x, batch_u = usets)

    def torch_dynamics_predict(self,cur_state:torch.tensor, batch_u = torch.tensor):
        #  x, y, psi, vx, vy, wz, z, roll, pitch

        assert len(batch_u.shape) ==3, 'uset is not in batch'
        batch_size = batch_u.shape[0]

        roll_state = cur_state.clone()
        # x_set = x_.repeat(n_sample,1,1)  #  [ number_of_sample, number_of_batch, number_of_states]

        pred_states = []
        pred_states.append(roll_state.clone())
        ####
        # N is the number of step to predict
        # same for nominal dynamics
        tmp_x_set_nominal = roll_state.clone()
        ## ## Update for each prediction step
        for i in range(self.N-1):
            px = tmp_x_set_nominal[:,0]
            py = tmp_x_set_nominal[:,1]
            yaw = tmp_x_set_nominal[:,2]
            vx = torch.max(tmp_x_set_nominal[:,3],torch.ones(len(tmp_x_set_nominal[:,3])).to(0)*0.05)
            vy = tmp_x_set_nominal[:,4]
            wz = tmp_x_set_nominal[:,5]
            z = tmp_x_set_nominal[:,6]
            usets = batch_u[:,i,:]
            axb = usets[:,0]
            delta = usets[:,1]

            if self.local_map.is_ready():
                pose = torch.transpose(torch.vstack([px,py,yaw]),0,1)
                rpy_tmp = self.local_map.get_rollpitch(pose)
                roll = rpy_tmp[:,0]
                pitch = rpy_tmp[:,1]
                tmp_x_set_nominal[:,6] = self.local_map.get_elevation(pose)
            else:
                roll = tmp_x_set_nominal[:,7]
                pitch = tmp_x_set_nominal[:,8]

            rot_base_to_world = b_to_g_rot(roll,pitch,yaw).double()

            Fzf, Fzr = self.compute_normal_force(tmp_x_set_nominal,usets,roll,pitch)
            alpha_f, alpha_r = self.compute_slip(tmp_x_set_nominal,usets)
            Fyf = Fzf * alpha_f
            Fyr =  Fzr * alpha_r
            local_vel = torch.hstack([vx.view(-1,1),vy.view(-1,1),torch.zeros(len(vx)).to(device=self.device).view(-1,1)]).view(-1,3,1).double()
            vel_in_world = torch.bmm(rot_base_to_world, local_vel).view(-1,3)

            vxw = vel_in_world[:,0]

            vyw = vel_in_world[:,1]

            vzw = vel_in_world[:,2]

            roll_state[:,0] += self.dt*vxw
            roll_state[:,1] += self.dt*vyw
            roll_state[:,2] += self.dt*(torch.cos(roll)/(torch.cos(pitch)+1e-10)*wz)
            roll_state[:,3] += self.dt*axb
            roll_state[:,3] = torch.max(torch.hstack((roll_state[:,3].view(-1,1),torch.zeros(len(roll_state[:,3])).to(device=self.device).view(-1,1))),dim =1).values
            roll_state[:,4] += self.dt*((Fyf+Fyr+self.m*self.g*torch.cos(pitch)*torch.sin(roll))/self.m-vx*wz)
            roll_state[:,5] += self.dt*((Fyf*self.Lf*torch.cos(delta)-self.Lr*Fyr)/self.Izz)
            roll_state[:,6] += self.dt*vzw
            pred_states.append(roll_state.clone())
            tmp_x_set_nominal = roll_state.clone()

        pred_states = torch.stack(pred_states,dim =1)
        xhat = self.get_batch_xhat(pred_states)

        return xhat, batch_u, pred_states

    def torch_2d_predict(self,cur_state:torch.tensor, batch_u = torch.tensor):
        #  x, y, psi, vx, vy, wz, z, roll, pitch

        assert len(batch_u.shape) ==3, 'uset is not in batch'
        batch_size = batch_u.shape[0]
        roll_state = cur_state.clone()

        pred_states = []
        pred_states.append(roll_state.clone())
        ####
        # N is the number of step to predict
        # same for nominal dynamics
        tmp_x_set_nominal = roll_state.clone()
        ## ## Update for each prediction step
        for i in range(self.N-1):
            px = tmp_x_set_nominal[:,0]
            py = tmp_x_set_nominal[:,1]
            yaw = tmp_x_set_nominal[:,2]
            vx = torch.max(tmp_x_set_nominal[:,3],torch.ones(len(tmp_x_set_nominal[:,3])).to(0)*0.05)
            vy = tmp_x_set_nominal[:,4]
            wz = tmp_x_set_nominal[:,5]
            z = tmp_x_set_nominal[:,6]
            usets = batch_u[:,i,:]
            axb = usets[:,0]
            delta = usets[:,1]

            if self.local_map.is_ready():
                pose = torch.transpose(torch.vstack([px,py,yaw]),0,1)
                rpy_tmp = self.local_map.get_rollpitch(pose)
                roll = rpy_tmp[:,0]
                pitch = rpy_tmp[:,1]
                tmp_x_set_nominal[:,6] = self.local_map.get_elevation(pose)
            else:
                roll = tmp_x_set_nominal[:,7]
                pitch = tmp_x_set_nominal[:,8]

            pitch = torch.zeros(pitch.shape).to(self.device)
            roll = torch.zeros(roll.shape).to(self.device)

            rot_base_to_world = b_to_g_rot(roll,pitch,yaw).double()

            Fzf, Fzr = self.compute_normal_force(tmp_x_set_nominal,usets,roll,pitch)
            alpha_f, alpha_r = self.compute_slip(tmp_x_set_nominal,usets)
            Fyf = Fzf * alpha_f
            Fyr =  Fzr * alpha_r
            local_vel = torch.hstack([vx.view(-1,1),vy.view(-1,1),torch.zeros(len(vx)).to(device=self.device).view(-1,1)]).view(-1,3,1).double()
            vel_in_world = torch.bmm(rot_base_to_world, local_vel).view(-1,3)

            vxw = vel_in_world[:,0]

            vyw = vel_in_world[:,1]

            roll_state[:,0] += self.dt*vxw
            roll_state[:,1] += self.dt*vyw
            roll_state[:,2] += self.dt*(torch.cos(roll)/(torch.cos(pitch)+1e-10)*wz)
            roll_state[:,3] += self.dt*axb
            roll_state[:,3] = torch.max(torch.hstack((roll_state[:,3].view(-1,1),torch.zeros(len(roll_state[:,3])).to(device=self.device).view(-1,1))),dim =1).values
            roll_state[:,4] += self.dt*((Fyf+Fyr+self.m*self.g*torch.cos(pitch)*torch.sin(roll))/self.m-vx*wz)
            roll_state[:,5] += self.dt*((Fyf*self.Lf*torch.cos(delta)-self.Lr*Fyr)/self.Izz)

            pred_states.append(roll_state.clone())
            tmp_x_set_nominal = roll_state.clone()

        pred_states = torch.stack(pred_states,dim =1)
        xhat = self.get_batch_xhat(pred_states)

        return xhat, batch_u, pred_states

    def get_batch_xhat(self, pred_states):
        # pred_states : x, y, psi, vx, vy, wz, z, roll, pitch
        # hxat : x, y, vx, vy, wz,  roll, pitch, yaw
        batch_size = pred_states.shape[0]
        xhat = torch.zeros(batch_size,pred_states.shape[1],8).to(self.device)
        xhat[:,:,0] = pred_states[:,:,0] # x -> x
        xhat[:,:,1] = pred_states[:,:,1] # y -> y
        xhat[:,:,2] = pred_states[:,:,3] # psi -> vx
        xhat[:,:,3] = pred_states[:,:,4] # vx -> vy
        xhat[:,:,4] = pred_states[:,:,5] # vy -> wz
        xhat[:,:,5] = pred_states[:,:,7] # wz -> roll
        xhat[:,:,6] = pred_states[:,:,8] # z -> pitch
        xhat[:,:,7] = pred_states[:,:,2] # roll -> yaw

        cur_states_x = xhat[:,0,0].clone()
        cur_states_y = xhat[:,0,1].clone()

        world_del_x = xhat[:,:,0].transpose(0,1) - cur_states_x.repeat(pred_states.shape[1],1)
        world_del_y = xhat[:,:,1].transpose(0,1) - cur_states_y.repeat(pred_states.shape[1],1)

        xhat[:,:,0] = world_del_x.transpose(0,1).clone() # delta px in local
        xhat[:,:,1] = world_del_y.transpose(0,1).clone() # delta py in local

        del_yaw = wrap_to_pi_torch(xhat[:,:,-1].transpose(0,1) - xhat[:,0,-1])

        xhat[:,:,-1] = del_yaw.transpose(0,1).clone()

        return xhat

    def action_sampling(self, cur_state: VehicleState):

        batch_num = 20
        batch_u = torch.zeros([batch_num,9,2]).to(self.device)

        ax_low_limit = self.ax_min
        ax_high_limit = self.ax_max
        batch_ax = torch.rand([batch_num,9]).to(self.device)
        batch_ax = (ax_high_limit - ax_low_limit) * batch_ax + ax_low_limit

        delta_low_imit = self.delta_min
        delta_high_imit = self.delta_max
        batch_delta = torch.rand([batch_num,9]).to(self.device)
        batch_delta = (delta_high_imit - delta_low_imit) * batch_delta + delta_low_imit

        batch_u[:,:,0] = batch_ax
        batch_u[:,:,1] = batch_delta

        for i in range(batch_num):
            batch_u[i,:,1] = -0.25+0.025*i

        odd_indices = torch.arange(1, batch_u.shape[0], step=2).to(self.device)
        even_indices = torch.arange(0, batch_u.shape[0], step=2).to(self.device)

        if cur_state.local_twist.linear.x < 1.0:
            batch_u[odd_indices,:,0] = 1.0
            batch_u[even_indices,:,0] = -1.0
        else:
            batch_u[odd_indices,:,0] = 1.0
            batch_u[even_indices,:,0] = -1.0

        return batch_u


from unittest import TestCase
import unittest
from equadratures import *
import numpy as np
from scipy.stats import skew, linregress, multivariate_normal

def fun(x):
    return 5.0 * x[0] ** 3 - x[0] * x[1] + 3.0 * x[1] * x[2] ** 3 + 32.0

class TestF(TestCase):
    def test_nataf(self):
        np.random.seed(1)
        zeta_1 = Parameter(distribution='truncated-gaussian', shape_parameter_A=3.0, shape_parameter_B=2.0,
                           order=15, lower=-2.0, upper=4.0)
        zeta_2 = Parameter(distribution='truncated-gaussian', shape_parameter_A=-1.0, shape_parameter_B=0.1,
                           order=15, lower=-5.0, upper=5.0)
        zeta_3 = Parameter(distribution='truncated-gaussian', shape_parameter_A=2.0, shape_parameter_B=2.0,
                           order=15, lower=0.0, upper=6.0)
        R = np.eye(3)
        R[0, 1] = 0.6
        R[0, 2] = 0.3
        R[2, 1] = 0.2
        R[1, 0] = R[0, 1]
        R[2, 0] = R[0, 2]
        R[1, 2] = R[2, 1]
        myBasis = Basis('tensor-grid')
        myPoly = Poly([zeta_1, zeta_2, zeta_3], myBasis, method='numerical-integration')
        myNataf = Correlations(R, poly=myPoly, method='nataf-transform')
        samples_mc = myNataf.get_correlated_samples(N=50000)
        f_mc = evaluate_model(samples_mc, fun)
        myNataf.set_model(fun)
        myTransformedPoly = myNataf.get_transformed_poly()
        mean, variance = myTransformedPoly.get_mean_and_variance()
        skewness, kurtosis = myTransformedPoly.get_skewness_and_kurtosis()

        np.testing.assert_almost_equal(mean / np.mean(f_mc), 1.0, decimal=1.5)
        np.testing.assert_almost_equal(variance / np.var(f_mc), 1.0, decimal=1.5)

        np.testing.assert_almost_equal(skewness / skew(f_mc)[0], 1.0, decimal=1.5)

    def test_GS(self):
        np.random.seed(1)
        zeta_1 = Parameter(distribution='truncated-gaussian', shape_parameter_A=3.0, shape_parameter_B=2.0,
                           order=5, lower=-2.0, upper=4.0)
        zeta_2 = Parameter(distribution='truncated-gaussian', shape_parameter_A=-1.0, shape_parameter_B=0.1,
                           order=5, lower=-5.0, upper=5.0)
        zeta_3 = Parameter(distribution='truncated-gaussian', shape_parameter_A=2.0, shape_parameter_B=2.0,
                           order=5, lower=0.0, upper=6.0)
        R = np.eye(3)
        R[0, 1] = 0.6
        R[0, 2] = 0.3
        R[2, 1] = 0.2
        R[1, 0] = R[0, 1]
        R[2, 0] = R[0, 2]
        R[1, 2] = R[2, 1]
        myBasis = Basis('tensor-grid')
        myPoly = Poly([zeta_1, zeta_2, zeta_3], myBasis, method='least-squares',
                      sampling_args={'mesh': 'monte-carlo', 'subsampling-algorithm': 'lu'})
        myGS = Correlations(R, poly=myPoly, method='gram-schmidt')
        samples_mc = myGS.get_correlated_samples(N=500)
        f_mc = evaluate_model(samples_mc, fun)
        myGS.set_model(fun)
        myTransformedPoly = myGS.get_transformed_poly()

        s, _, r, _, _ = linregress(myTransformedPoly.get_polyfit(samples_mc).reshape(-1),
                                   f_mc.reshape(-1))
        np.testing.assert_almost_equal(s, 1.0, decimal=2)
        np.testing.assert_almost_equal(r, 1.0, decimal=2)

    def test_pdf(self):
        np.random.seed(1)
        X_test = np.random.uniform(-3, 3, (1000, 2))
        my_params = [Parameter(distribution='gaussian', shape_parameter_A=0, shape_parameter_B=1.0,
                               order=3) for _ in range(2)]
        corr_mat = np.eye(2)
        corr_mat[0,1] = corr_mat[1,0] = 0.5
        my_corr = Correlations(corr_mat, parameters=my_params)
        test_pdf = my_corr.get_pdf(X_test)
        truth_pdf = multivariate_normal.pdf(X_test, mean=[0,0], cov=corr_mat)
        np.testing.assert_almost_equal(np.linalg.norm(test_pdf - truth_pdf), 0.0, decimal=2)

if __name__ == '__main__':
    unittest.main()


#!/usr/bin/env python

#  This script tells if a File, IP, Domain or URL may be malicious according to the data in OTX
import IndicatorTypes
import json
import geoip2.database
import socket
import requests
import os
import time
import logging

from datetime import datetime
from dateutil.parser import parse
from OTXv2 import OTXv2

# Define logger for cross appliction logging consistency
logger = logging.getLogger(__name__)

def getValue(results, keys):
    '''
    Get a nested key from a dict, without having to do loads of ifs
    '''
    if type(keys) is list and len(keys) > 0:

        if type(results) is dict:
            key = keys.pop(0)
            if key in results:
                return getValue(results[key], keys)
            else:
                return None
        else:
            if type(results) is list and len(results) > 0:
                return getValue(results[0], keys)
            else:
                return results
    else:
        return results

def ip(otx, ip, mmdb, query_time):
    '''
    Query AlienVault OTX for malicious IP checking
    '''
    alerts = {}
    report_db = ''
    latitude = ''
    longitude = ''
    full_url_list = []
    top_five_urls = [] # List of 5 most recent urls
    url_set = set() # Remove duplicates from full_url_list
    first_reported = ''
    last_reported = ''
    url_status = ''

    try:
        result = otx.get_indicator_details_full(IndicatorTypes.IPv4, ip)
        pulses = getValue(result['general'], ['pulse_info', 'pulses'])
        if pulses:
            report_db = 'otx.alienvault.com'
            latitude = result['geo']['latitude']
            longitude = result['geo']['longitude']
            full_url_list = getValue(result['passive_dns'], ['passive_dns'])

            for url in full_url_list:
                if 'hostname' and 'flag_url' in url:
                    url_set.add(url['hostname']+'/'+url['flag_url'])

            top_five_urls = list(url_set)[:5]
            first_reported = parse(full_url_list[-1]['first']).isoformat()
            last_reported = parse(full_url_list[0]['last']).isoformat()
            url_status = 'potentially_malicious'
        else:
            response = geoip2.database.Reader(mmdb).city(ip)
            report_db = 'otx.alienvault.com'
            latitude = response.location.latitude
            longitude = response.location.longitude
            try:
                top_five_urls = [socket.gethostbyaddr(ip)[0]]
            except socket.herror:
                top_five_urls = []

            first_reported = ''
            last_reported = ''
            url_status = 'likely_benign'

    except Exception as e:
        logger.exception(" :  You received this error with the OTX API Data... {}".format(e))

    # Build alerts dictionary from variables
    alerts = {'report_db': report_db, 'ip_addr': ip, 'latitude': latitude, 'longitude': longitude,
              'urls': top_five_urls, 'first_reported': first_reported, 'last_reported': last_reported,
              'query_date': query_time, 'url_status': url_status}

    return alerts

def hostname(host, mmdb, query_url, query_time):
    '''
    Query urlhaus.abuse.ch for malicious host checking
    '''
    alerts = {}
    ip_addr = ''
    report_db = ''
    latitude = ''
    longitude = ''
    full_url_list = []
    url_set = set() # Remove duplicates from full_url_list
    first_reported = ''
    last_reported = ''
    url_status = ''
    __version__ = '0.0.2'

    try:
        r = requests.post("{}host/".format(query_url), headers={"User-Agent" : "urlhaus-python-client-{}".format(__version__)}, data={"host": host})
        if r.ok:
            geo_reader = geoip2.database.Reader(mmdb)
            if r.json()['query_status'] == "no_results":
                report_db = 'urlhaus.abuse.ch'
                try:
                    ip_addr = socket.gethostbyname(host)
                except socket.error:
                    ip_addr = ''

                if ip_addr:
                    try:
                        latitude = geo_reader.city(ip_addr).location.latitude
                        longitude = geo_reader.city(ip_addr).location.longitude
                    except:
                        latitude = ''
                        longitude = ''

                full_url_list = [host]
                first_reported = ''
                last_reported = ''
                url_status = 'likely_benign'

            else:
                report_db = 'urlhaus.abuse.ch'
                try:
                    ip_addr = socket.gethostbyname(host)
                except socket.error:
                    ip_addr = ''

                if ip_addr:
                    try:
                        latitude = geo_reader.city(ip_addr).location.latitude
                        longitude = geo_reader.city(ip_addr).location.longitude
                    except:
                        latitude = ''
                        longitude = ''

                for url in r.json()['urls']:
                    url_set.add(url['url'].split('//', 2)[1])

                full_url_list = list(url_set)[:5]
                first_reported = parse(r.json()['firstseen']).isoformat()
                if r.json()['urls'][0]['url_status'] == 'online':
                    url_status = 'potentially_malicious'
                else:
                    url_status = 'potentially_malicious_but_offline'

        else:
            logger.error(" :  Unable to read response as json")

    except Exception as e:
        logger.exception(" :  Unable to connect to URLHaus API. Recieved the following error {}".format(e))

    # Build alerts dictionary from variables
    alerts = {'report_db': report_db, 'ip_addr': ip_addr, 'latitude': latitude, 'longitude': longitude,
              'urls': full_url_list, 'first_reported': first_reported, 'last_reported': last_reported,
              'query_date': query_time, 'url_status': url_status}
    return alerts

def allow_list():
    '''
    Check for existence of domain in the top 1 million domains list and don't lookup
    the value if it is. This will miss the instances where a bad actor uses Google, Twitter,
    Dropbox, etc. for C2, but it significantly reduces network overhead and processing overhead
    from continual lookups of "known good" URLs. These URLs will still be analyzed by the
    tls-mal-detect.py file, so it is not a total loss.
    '''

def sslbl():
    '''
    Check IP against SSL Blocklist using sslbl.abuse.ch
    '''

def main():
    '''
    Run some automated tests for ip and hostname methods
    '''
    api_key = os.environ.get('API_KEY')
    otx_server = 'https://otx.alienvault.com/'
    otx = OTXv2(api_key, server=otx_server)

    urlhaus_api = "https://urlhaus-api.abuse.ch/v1/"

    query_timestamp = datetime.utcnow().isoformat()
    mmdb_path = '../netcap/resources/GeoLite2-City.mmdb'

    ip_addr = '4.4.4.4'
    host = 'google.com'

    print("IP Query output: {}".format(ip(otx, ip_addr, mmdb_path, query_timestamp)))
    print("Hostname Query output: {}".format(hostname(host, mmdb_path, urlhaus_api, query_timestamp)))

if __name__ == "__main__":
    try:
        exit(main())
    except Exception:
        logging.exception("Exception in main()")
        exit(1)