<commit_before># separate a folder into train,validate and test set and dump thus into three file
from __future__ import division
from random import shuffle
import pickle
import fnmatch
import os
import codecs
from net.wyun.mer.ink.sample import Sample
from os.path import basename
from scipy import misc
from xlsx import Excel
from net.wyun.mer.ink.scg import Scg

'''
   from a folder contains inkml files (files ending with .inkml), generate 3 list
   one for dev, one for training, and one for validate. the program will get inkml files
   recursively.
   the *.flist file is an pickle file which contains list of file names (absolute path version)

   usage: test_batch.py for more information
'''

import logging
logger= logging.getLogger("batch")
logging.basicConfig(format='%(asctime)s - %(funcName)s - %(message)s', filename='temp/batch.log',level=logging.DEBUG)

class Batch(object):
    def __init__(self, inkml_file_path, xlsx_file_path, list_out_dir='data/batch/pickle/'):
        '''

        :param inkml_file_path: directory contains *.inkml files
        :param list_out_dir: 3 list would be generated, one for training, one for test, and the other validation
        '''
        self.inkml_file_path = os.path.realpath(inkml_file_path)  # absolute path
        self.xlsx_file_path = os.path.realpath(xlsx_file_path)
        self.list_out_dir = list_out_dir
        self.latex_index = 0
        self.latex_length_threshold = 6 # if length < 4, this inkml file is not included
        self.scgs = self.load_all_scgs()
        self.id_blacklist = self.blacklist()

    def blacklist(self):
        l = []
        l.append(1509259382091)
        l.append(1509239879109)
        return l

    def generate_scg_idlists(self):
        '''
        handle xlsx files
        from a list of xlsx files, generate three scg list: train, validate, and test.
        it store the ids for scg files embedded in the excel files
        :return:
        '''
        ids = list(self.scgs.keys())
        ids_2nd = []

        for id in ids:
            if id in self.id_blacklist:
                continue

            scg = self.scgs[id]
            if scg.get_latex() == '1' or scg.get_latex() == 'i':
                continue

            aspect = scg.w_h_ratio
            if aspect > 80.0 or aspect < 0.02: #scg.get_latex() == '1':
                print '==========> scg aspect: ', aspect
                continue
            scg.save_image('data/batch/formula_images/' + str(id) + '.png')
            ids_2nd.append(id)
        ids = ids_2nd

        shuffle(ids)
        print 'total ids: ', len(ids)

        total = len(ids)
        validate_num = int(total * 10 / 100)
        test_num = validate_num

        print 'validate num(ids): ', validate_num
        validate = ids[:validate_num]

        test = ids[validate_num:validate_num + test_num]
        train = ids[validate_num + test_num:]

        print 'test num(ids): ', len(test)
        print 'train num(ids): ', len(train)

        list_out_dir = self.list_out_dir
        with open(list_out_dir + 'validate.idlist', 'wb') as fout:
            pickle.dump(validate, fout)

        with open(list_out_dir + 'test.idlist', 'wb') as fout:
            pickle.dump(test, fout)

        with open(list_out_dir + 'train.idlist', 'wb') as fout:
            pickle.dump(train, fout)

        with open(list_out_dir + 'total.idlist', 'wb') as fout:
            pickle.dump(ids, fout)

    def load_all_scgs(self):
        '''
        From a list of .xlsx files, generate a dictionary of Scgs with key is the id of Scg
        :return: scgs, dictionary of Scgs
        '''
        scgs = {}  # dictionary of Scgs
        aspect_ratios = {}
        with codecs.open('data/invalid_scgs.txt', 'w', 'utf-8') as f_out:
            for root, dirnames, filenames in os.walk(self.xlsx_file_path, followlinks=True):
                for filename in fnmatch.filter(filenames, '*.xlsx'):
                    tmp_path = os.path.join(root, filename)
                    # print tmp_path
                    excel = Excel(tmp_path)
                    max_row = excel.get_row_number()
                    for idx in range(1, max_row):
                        print idx
                        id, scg, truth = excel.get_scg_record(idx)
                        if len(truth.strip()) == 0:
                            print id
                            line = str(id) + ' ' + 'no latex\n'
                            f_out.write(line)
                            continue
                        scgs[id] = Scg(id, scg, truth)
                        aspect_ratios[id] = scgs[id].w_h_ratio
        #dump aspect_ratios to a file
        ars = aspect_ratios.values()
        print 'aspect ratio min, mean, max: ', min(ars), max(ars), sum(ars) / float(len(ars))

        with open(self.list_out_dir + 'aspectratio.dict', 'wb') as fout:
            pickle.dump(aspect_ratios, fout)
        return scgs

    def generate_filelists(self):
        '''
        retrieve all inkml files in self.inkml_file_path, check if it has latex truth information. if yes, add it to
        a list files. divide files into 3 categories. One for training, one for test, and the other validate. the share
        is in 80:10:10.
        :return:
        '''
        files = []
        files_nolatex = [] # it also includes those  length < 4
        for root, dirnames, filenames in os.walk(self.inkml_file_path, followlinks=True):
            for filename in fnmatch.filter(filenames, '*.xlsx'):
                tmp_path = os.path.join(root, filename)
                #print tmp_path
                sample = Sample(tmp_path)
                if hasattr(sample, 'latex') and self.check_latex_length(sample):  #latex has been stripped
                    #print 'latex: ', tmp_path

                    files.append(tmp_path)
                else:
                    files_nolatex.append(tmp_path)

        shuffle(files)
        print 'total files: ', len(files)

        total = len(files)
        validate_num = int(total * 10 / 100)
        test_num = validate_num

        print 'validate num: ', validate_num
        validate = files[:validate_num]

        test = files[validate_num:validate_num + test_num]
        train = files[validate_num + test_num:]

        print 'test num: ', len(test)
        print 'train num: ', len(train)
        print 'files w/o latex num: ', len(files_nolatex)

        list_out_dir = self.list_out_dir
        with open(list_out_dir + 'validate.flist', 'wb') as fout:
            pickle.dump(validate, fout)

        with open(list_out_dir + 'test.flist', 'wb') as fout:
            pickle.dump(test, fout)

        with open(list_out_dir + 'train.flist', 'wb') as fout:
            pickle.dump(train, fout)

        with open(list_out_dir + 'total.flist', 'wb') as fout:
            pickle.dump(files, fout)

        with open(list_out_dir + 'nolatex.flist', 'wb') as fout:
            pickle.dump(files_nolatex, fout)

    def check_latex_length(self, sample):
        return True
        #return len(sample.latex) >= self.latex_length_threshold

    def process_list(self, pair_file_name, latex_out):
        '''

        :param pair_file_name, set name: for example, 'train', one train.flist, the other train.idlist
        :param outfile: for example, data/batch/train.lst
        :param latex_out: 'data/batch/latex_list.txt'
        :return:
        '''
        print "No dir specified, using default dir"
        base_dir = 'data/batch/'  # self.list_out_dir
        flist_base_dir = self.list_out_dir
        prefix = 'im2latex_'

        infile = flist_base_dir + pair_file_name + '.flist'     # data/batch/pickle/train.flist
        outfile = base_dir + prefix + pair_file_name + '.lst'   # data/batch/im2latex_train.lst
        scg_id_list_file = flist_base_dir + pair_file_name + '.idlist'

        print 'input: ', infile
        print 'output: ', outfile

        with open(infile, 'rb') as fin:
            flist = pickle.load(fin)

        with open(scg_id_list_file, 'rb') as fin:
            idlist = pickle.load(fin)

        png_dir = 'data/batch/formula_images/'
        with codecs.open(outfile, 'w', 'utf-8') as f_out:
            for file_path in flist:
                print 'file: ', file_path
                # process one file
                sample = Sample(file_path)
                latex = sample.latex
                #output latex
                latex_out.write(latex[1:-1] + '\n')
                #get filename without extension
                png_name = self.get_filename_noext(file_path)
                line = str(self.latex_index) + ' ' + png_name + ' ' + 'basic\n'
                f_out.write(line)
                #output image
                img, W, H = sample.render()
                misc.imsave(png_dir + png_name + '.png', img)
                self.latex_index += 1
            for id in idlist:
                print 'id: ', id
                logging.info('processing %d scg record', id)
                scg = self.scgs[id]
                latex = scg.get_latex()
                # output latex
                latex_out.write(latex + '\n')
                # get filename without extension
                png_name = str(id)
                line = str(self.latex_index) + ' ' + png_name + ' ' + 'basic\n'
                f_out.write(line)
                # output image (this task is done when dividing the list into 3 groups
                #scg.save_image(png_dir + png_name + '.png')
                self.latex_index += 1

    def get_filename_noext(self, file_path):
        # now you can call it directly with basename
        base = basename(file_path)
        return os.path.splitext(base)[0]

    def generate_im2markup_files(self):
        '''
        process train.flist, test.flist, and validate.flist
        produce train.lst, test.lst, and validate.lst
        :return:
        '''
        #with open("latex_list.txt", "w") as f:
        with codecs.open('data/batch/im2latex_formulas.lst', 'w', 'utf-8') as latex_out:
            #line = str(mydict[s]) + ' ' + s[0] + ' ' + str(s[1])
            #       fout.write(line + '\n')
            print "No dir specified, using default dir"
            base_dir = 'data/batch/' #self.list_out_dir
            flist_base_dir = self.list_out_dir
            prefix = 'im2latex_'
            pair_file_name = 'train'
            print 'generate ', pair_file_name, ' lst'
            self.process_list(pair_file_name, latex_out)
            pair_file_name = 'validate'
            print 'generate ', pair_file_name, ' lst'
            self.process_list(pair_file_name, latex_out)
            pair_file_name = 'test'
            print 'generate ', pair_file_name, ' lst'
            self.process_list(pair_file_name, latex_out)
<commit_msg>1. png_dir now a module level variable. it is shared in other places. 2. add inkml2png function which generates png images for inkml files. 3. because scgimage class is used to generate png images for scg contents need to make sure to use the same Height for the generated png images.<commit_after># separate a folder into train,validate and test set and dump thus into three file
from __future__ import division
from random import shuffle
import pickle
import fnmatch
import os
import codecs
from net.wyun.mer.ink.sample import Sample
from os.path import basename
from scipy import misc
from xlsx import Excel
from net.wyun.mer.ink.scg import Scg

'''
   from a folder contains inkml files (files ending with .inkml), generate 3 list
   one for dev, one for training, and one for validate. the program will get inkml files
   recursively.
   the *.flist file is an pickle file which contains list of file names (absolute path version)

   usage: test_batch.py for more information
'''

import logging
logger= logging.getLogger("batch")
logging.basicConfig(format='%(asctime)s - %(funcName)s - %(message)s', filename='temp/batch.log',level=logging.DEBUG)

png_dir = 'data/batch/formula_images/'
class Batch(object):
    def __init__(self, inkml_file_path, xlsx_file_path, list_out_dir='data/batch/pickle/'):
        '''

        :param inkml_file_path: directory contains *.inkml files
        :param list_out_dir: 3 list would be generated, one for training, one for test, and the other validation
        '''
        self.inkml_file_path = os.path.realpath(inkml_file_path)  # absolute path
        self.xlsx_file_path = os.path.realpath(xlsx_file_path)
        self.list_out_dir = list_out_dir
        self.latex_index = 0
        self.latex_length_threshold = 6 # if length < 4, this inkml file is not included
        #self.scgs = self.load_all_scgs()
        self.id_blacklist = self.blacklist()

    def blacklist(self):
        l = []
        l.append(1509259382091)
        l.append(1509239879109)
        return l

    def generate_scg_idlists(self):
        '''
        handle xlsx files
        from a list of xlsx files, generate three scg list: train, validate, and test.
        it store the ids for scg files embedded in the excel files
        :return:
        '''
        ids = list(self.scgs.keys())
        ids_2nd = []

        for id in ids:
            if id in self.id_blacklist:
                continue

            scg = self.scgs[id]
            if scg.get_latex() == '1' or scg.get_latex() == 'i':
                continue

            aspect = scg.w_h_ratio
            if aspect > 80.0 or aspect < 0.02: #scg.get_latex() == '1':
                print '==========> scg aspect: ', aspect
                continue
            scg.save_image('data/batch/formula_images/' + str(id) + '.png')
            ids_2nd.append(id)
        ids = ids_2nd

        shuffle(ids)
        print 'total ids: ', len(ids)

        total = len(ids)
        validate_num = int(total * 10 / 100)
        test_num = validate_num

        print 'validate num(ids): ', validate_num
        validate = ids[:validate_num]

        test = ids[validate_num:validate_num + test_num]
        train = ids[validate_num + test_num:]

        print 'test num(ids): ', len(test)
        print 'train num(ids): ', len(train)

        list_out_dir = self.list_out_dir
        with open(list_out_dir + 'validate.idlist', 'wb') as fout:
            pickle.dump(validate, fout)

        with open(list_out_dir + 'test.idlist', 'wb') as fout:
            pickle.dump(test, fout)

        with open(list_out_dir + 'train.idlist', 'wb') as fout:
            pickle.dump(train, fout)

        with open(list_out_dir + 'total.idlist', 'wb') as fout:
            pickle.dump(ids, fout)

    def load_all_scgs(self):
        '''
        From a list of .xlsx files, generate a dictionary of Scgs with key is the id of Scg
        :return: scgs, dictionary of Scgs
        '''
        scgs = {}  # dictionary of Scgs
        aspect_ratios = {}
        with codecs.open('data/invalid_scgs.txt', 'w', 'utf-8') as f_out:
            for root, dirnames, filenames in os.walk(self.xlsx_file_path, followlinks=True):
                for filename in fnmatch.filter(filenames, '*.xlsx'):
                    tmp_path = os.path.join(root, filename)
                    # print tmp_path
                    excel = Excel(tmp_path)
                    max_row = excel.get_row_number()
                    for idx in range(1, max_row):
                        print idx
                        id, scg, truth, req_at, resp_at = excel.get_scg_record(idx)
                        if len(truth.strip()) == 0:
                            print id
                            line = str(id) + ' ' + 'no latex\n'
                            f_out.write(line)
                            continue
                        scgs[id] = Scg(id, scg, truth, req_at, resp_at)
                        aspect_ratios[id] = scgs[id].w_h_ratio
        #dump aspect_ratios to a file
        ars = aspect_ratios.values()
        print 'aspect ratio min, mean, max: ', min(ars), max(ars), sum(ars) / float(len(ars))

        with open(self.list_out_dir + 'aspectratio.dict', 'wb') as fout:
            pickle.dump(aspect_ratios, fout)
        return scgs

    def all_inkml_files(self):
        '''

        :return: files, files_nolatex
        '''
        files = []
        files_nolatex = []
        for root, dirnames, filenames in os.walk(self.inkml_file_path, followlinks=True):
                for filename in fnmatch.filter(filenames, '*.inkml'):
                    tmp_path = os.path.join(root, filename)
                    print tmp_path
                    sample = Sample(tmp_path)
                    if hasattr(sample, 'latex') and self.check_latex_length(sample):  #latex has been stripped
                        #print 'latex: ', tmp_path

                        files.append(tmp_path)
                    else:
                        files_nolatex.append(tmp_path)
        return files, files_nolatex

    def generate_filelists(self):
        '''
        retrieve all inkml files in self.inkml_file_path, check if it has latex truth information. if yes, add it to
        a list files. divide files into 3 categories. One for training, one for test, and the other validate. the share
        is in 80:10:10.
        :return:
        '''
        files, files_nolatex = self.all_inkml_files()

        shuffle(files)
        print 'total files: ', len(files)

        total = len(files)
        validate_num = int(total * 10 / 100)
        test_num = validate_num

        print 'validate num: ', validate_num
        validate = files[:validate_num]

        test = files[validate_num:validate_num + test_num]
        train = files[validate_num + test_num:]

        print 'test num: ', len(test)
        print 'train num: ', len(train)
        print 'files w/o latex num: ', len(files_nolatex)

        list_out_dir = self.list_out_dir
        with open(list_out_dir + 'validate.flist', 'wb') as fout:
            pickle.dump(validate, fout)

        with open(list_out_dir + 'test.flist', 'wb') as fout:
            pickle.dump(test, fout)

        with open(list_out_dir + 'train.flist', 'wb') as fout:
            pickle.dump(train, fout)

        with open(list_out_dir + 'total.flist', 'wb') as fout:
            pickle.dump(files, fout)

        with open(list_out_dir + 'nolatex.flist', 'wb') as fout:
            pickle.dump(files_nolatex, fout)

    def check_latex_length(self, sample):
        return True
        #return len(sample.latex) >= self.latex_length_threshold

    def process_list(self, pair_file_name, latex_out):
        '''

        :param pair_file_name, set name: for example, 'train', one train.flist, the other train.idlist
        :param outfile: for example, data/batch/train.lst
        :param latex_out: 'data/batch/latex_list.txt'
        :return:
        '''
        print "No dir specified, using default dir"
        base_dir = 'data/batch/'  # self.list_out_dir
        flist_base_dir = self.list_out_dir
        prefix = 'im2latex_'

        infile = flist_base_dir + pair_file_name + '.flist'     # data/batch/pickle/train.flist
        outfile = base_dir + prefix + pair_file_name + '.lst'   # data/batch/im2latex_train.lst
        scg_id_list_file = flist_base_dir + pair_file_name + '.idlist'

        print 'input: ', infile
        print 'output: ', outfile

        with open(infile, 'rb') as fin:
            flist = pickle.load(fin)

        with open(scg_id_list_file, 'rb') as fin:
            idlist = pickle.load(fin)

        with codecs.open(outfile, 'w', 'utf-8') as f_out:
            for file_path in flist:
                print 'file: ', file_path
                # process one file
                sample = Sample(file_path)
                latex = sample.latex
                #output latex
                latex_out.write(latex[1:-1] + '\n')
                #get filename without extension
                png_name = self.get_filename_noext(file_path)
                line = str(self.latex_index) + ' ' + png_name + ' ' + 'basic\n'
                f_out.write(line)
                #output image
                img, W, H = sample.render()
                misc.imsave(png_dir + png_name + '.png', img)
                self.latex_index += 1
            for id in idlist:
                print 'id: ', id
                logging.info('processing %d scg record', id)
                scg = self.scgs[id]
                latex = scg.get_latex()
                # output latex
                latex_out.write(latex + '\n')
                # get filename without extension
                png_name = str(id)
                line = str(self.latex_index) + ' ' + png_name + ' ' + 'basic\n'
                f_out.write(line)
                # output image (this task is done when dividing the list into 3 groups
                #scg.save_image(png_dir + png_name + '.png')
                self.latex_index += 1

    def get_filename_noext(self, file_path):
        # now you can call it directly with basename
        base = basename(file_path)
        return os.path.splitext(base)[0]

    # This function is deprecated. Now all im2markup files would be generated based on database equation table
    def generate_im2markup_files(self):
        '''
        process train.flist, test.flist, and validate.flist
        produce train.lst, test.lst, and validate.lst
        :return:
        '''
        #with open("latex_list.txt", "w") as f:
        with codecs.open('data/batch/im2latex_formulas.lst', 'w', 'utf-8') as latex_out:
            #line = str(mydict[s]) + ' ' + s[0] + ' ' + str(s[1])
            #       fout.write(line + '\n')
            print "No dir specified, using default dir"
            base_dir = 'data/batch/' #self.list_out_dir
            flist_base_dir = self.list_out_dir
            prefix = 'im2latex_'
            pair_file_name = 'train'
            print 'generate ', pair_file_name, ' lst'
            self.process_list(pair_file_name, latex_out)
            pair_file_name = 'validate'
            print 'generate ', pair_file_name, ' lst'
            self.process_list(pair_file_name, latex_out)
            pair_file_name = 'test'
            print 'generate ', pair_file_name, ' lst'
            self.process_list(pair_file_name, latex_out)

    def inkml2png(self):
        '''
        for all inkml files in data/batch/inkml folder, generate its png image and save it
        to data/batch/formula_images/ folder

        It is used if different resolution images need to be generated for mhr performance evaluation.
        :return:
        '''
        files, files_nolatex = self.all_inkml_files()
        for file_path in files:
            print 'file: ', file_path
            # process one file
            sample = Sample(file_path)
            # get filename without extension
            png_name = self.get_filename_noext(file_path)
            # output image
            img, W, H = sample.render()
            misc.imsave(png_dir + png_name + '.png', img)
<|endoftext|>


<commit_before>__version__ = '3.2'
<commit_msg>bump webrecorder version to 3.3<commit_after>__version__ = '3.3'
<|endoftext|>


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

# Copyright (C) 2011 Adriano Monteiro Marques
#
# Author: Piotrek Wasilewski <wasilewski.piotrek@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import datetime

from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.contrib.contenttypes.models import ContentType
from django.core.urlresolvers import reverse
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger
from django.views.generic.create_update import create_object, update_object, \
    delete_object
from django.views.generic.list_detail import object_list, object_detail
from django.views.generic.simple import direct_to_template, redirect_to
from django.http import Http404
from search.core import search

from netadmin.events.models import Event
from netadmin.permissions.utils import filter_user_objects, \
    get_object_or_forbidden, grant_access, grant_edit, revoke_access, \
    revoke_edit, user_has_access

from models import Host, Network, NetworkHost
from forms import HostCreateForm, HostUpdateForm, NetworkCreateForm, \
    NetworkUpdateForm

@login_required
def host_list(request, page=None):
    search_phrase = request.GET.get('s')
    if search_phrase:
        hosts = search(Host, search_phrase)
    else:
        hosts = filter_user_objects(request.user, Host)

    paginator = Paginator(list(hosts), 10)

    page = page or request.GET.get('page', 1)
    try:
        hosts = paginator.page(page)
    except PageNotAnInteger:
        hosts = paginator.page(1)
    except EmptyPage:
        hosts = paginator.page(paginator.num_pages)

    extra_context = {
        'hosts': hosts,
        'url': reverse('host_list')
    }
    return direct_to_template(request, 'networks/host_list.html',
                              extra_context=extra_context)

@login_required
def host_detail(request, object_id):

    host, edit = get_object_or_forbidden(Host, object_id, request.user)

    extra_context = {
        'can_edit': edit
    }

    return object_detail(request, Host.objects.all(), object_id,
                         extra_context=extra_context)

@login_required
def host_create(request):

    if request.method == 'POST':
        form = HostCreateForm(request.POST)
        if form.is_valid():
            host = form.save()
            return redirect_to(request, url=host.get_absolute_url())

    extra_context = {
        'form': HostCreateForm(initial={'user': request.user.pk})
    }
    return direct_to_template(request, 'networks/host_form.html', extra_context)

@login_required
def host_update(request, object_id):

    host, edit = get_object_or_forbidden(Host, object_id, request.user)

    if not edit:
        raise Http404()

    return update_object(request, object_id=object_id,
                         form_class=HostUpdateForm,
                         template_name='networks/host_update.html')

@login_required
def host_delete(request, object_id):

    host, edit = get_object_or_forbidden(Host, object_id, request.user)

    if host.user != request.user:
        raise Http404()

    return delete_object(request, object_id=object_id, model=Host,
                         post_delete_redirect=reverse('host_list'))

@login_required
def network_list(request, page=None):
    search_phrase = request.GET.get('s')
    if search_phrase:
        nets = search(Network, search_phrase)
    else:
        nets = filter_user_objects(request.user, Network)

    paginator = Paginator(list(nets), 10)

    page = page or request.GET.get('page', 1)
    try:
        nets = paginator.page(page)
    except PageNotAnInteger:
        nets = paginator.page(1)
    except EmptyPage:
        nets = paginator.page(paginator.num_pages)

    extra_context = {
        'networks': nets,
        'url': '/network/network/list/'
    }
    return direct_to_template(request, 'networks/network_list.html',
                              extra_context=extra_context)

@login_required
def network_detail(request, object_id):
    """
    Network detail page has the following features:
        * displaying basic network info (name, description, etc.)
        * listing hosts related to network
        * creating relations between network and host
        * removing relations between network and host(s)
    """
    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    # remove relation between the network and selected host(s)
    if request.POST.getlist('remove_host'):
        if edit:
            hosts_pk = request.POST.getlist('remove_host')
            network_host = NetworkHost.objects.filter(network=network,
                                                      host__pk__in=hosts_pk)
            network_host.delete()

    # create relation between the network and selected host
    if request.POST.get('add_host'):
        if edit:
            host = Host.objects.get(pk=request.POST.get('add_host'))
            network_host = NetworkHost(network=network, host=host)
            network_host.save()

    queryset = Network.objects.filter(user=request.user)
    if network.hosts():
        hosts_ids = [host.pk for host in network.hosts()]
        hosts_other = Host.objects.exclude(pk__in=hosts_ids).filter(user=request.user)
    else:
        hosts_other = Host.objects.filter(user=request.user)
    extra_context = {
        'hosts_other': hosts_other,
        'can_edit': edit
    }
    return object_detail(request, queryset, object_id,
                         extra_context=extra_context)

@login_required
def network_create(request):

    if request.method == 'POST':
        form = NetworkCreateForm(request.POST)
        if form.is_valid():
            network = form.save()
            return redirect_to(request, url=network.get_absolute_url())

    extra_context = {
        'form': NetworkCreateForm(initial={'user': request.user.pk})
    }
    return direct_to_template(request, 'networks/network_form.html',
                              extra_context)

@login_required
def network_update(request, object_id):

    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    if not edit:
        raise Http404()

    return update_object(request, object_id=object_id,
                         form_class=NetworkUpdateForm,
                         template_name='networks/network_update.html')

@login_required
def network_delete(request, object_id):

    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    if network.user != request.user:
        raise Http404()

    return delete_object(request, object_id=object_id, model=Network,
                         post_delete_redirect=reverse('network_list'))

@login_required
def network_events(request, object_id):
    """Display events related to network"""

    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    queryset = Network.objects.all()
    related_hosts = [nh.host.pk for nh in NetworkHost.objects.filter(network=network)]
    events = Event.objects.filter(source_host__pk__in=related_hosts)
    extra_context = {
        'events': events,
        'can_edit': edit
    }
    return object_detail(request, queryset, object_id,
                         extra_context=extra_context,
                         template_name='networks/network_events.html')

@login_required
def share(request, object_type, object_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    user_id = request.POST.get('share')
    if user_id:
        user = User.objects.get(pk=user_id)
        grant_access(obj, user)
    return share_list(request, object_type, object_id)

@login_required
def share_not(request, object_type, object_id, user_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    user = User.objects.get(pk=user_id)
    revoke_access(obj, user)
    return share_list(request, object_type, object_id)

@login_required
def share_edit(request, object_type, object_id, user_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    user = User.objects.get(pk=user_id)
    if edit:
        revoke_edit(obj, user)
    else:
        grant_edit(obj, user)
    return share_list(request, object_type, object_id)

@login_required
def share_list(request, object_type, object_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    all_users = User.objects.exclude(pk=request.user.pk)
    other_users = []
    for user in all_users:
        if not user_has_access(obj, user):
            other_users.append(user)
    extra_context = {
        'object': obj,
        'object_type': object_type,
        'other_users': other_users
    }
    return direct_to_template(request, 'networks/share.html', extra_context)
<commit_msg>Removed unused import from networks.views<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (C) 2011 Adriano Monteiro Marques
#
# Author: Piotrek Wasilewski <wasilewski.piotrek@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import datetime

from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.core.urlresolvers import reverse
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger
from django.views.generic.create_update import create_object, update_object, \
    delete_object
from django.views.generic.list_detail import object_list, object_detail
from django.views.generic.simple import direct_to_template, redirect_to
from django.http import Http404
from search.core import search

from netadmin.events.models import Event
from netadmin.permissions.utils import filter_user_objects, \
    get_object_or_forbidden, grant_access, grant_edit, revoke_access, \
    revoke_edit, user_has_access

from models import Host, Network, NetworkHost
from forms import HostCreateForm, HostUpdateForm, NetworkCreateForm, \
    NetworkUpdateForm

@login_required
def host_list(request, page=None):
    search_phrase = request.GET.get('s')
    if search_phrase:
        hosts = search(Host, search_phrase)
    else:
        hosts = filter_user_objects(request.user, Host)

    paginator = Paginator(list(hosts), 10)

    page = page or request.GET.get('page', 1)
    try:
        hosts = paginator.page(page)
    except PageNotAnInteger:
        hosts = paginator.page(1)
    except EmptyPage:
        hosts = paginator.page(paginator.num_pages)

    extra_context = {
        'hosts': hosts,
        'url': reverse('host_list')
    }
    return direct_to_template(request, 'networks/host_list.html',
                              extra_context=extra_context)

@login_required
def host_detail(request, object_id):

    host, edit = get_object_or_forbidden(Host, object_id, request.user)

    extra_context = {
        'can_edit': edit
    }

    return object_detail(request, Host.objects.all(), object_id,
                         extra_context=extra_context)

@login_required
def host_create(request):

    if request.method == 'POST':
        form = HostCreateForm(request.POST)
        if form.is_valid():
            host = form.save()
            return redirect_to(request, url=host.get_absolute_url())

    extra_context = {
        'form': HostCreateForm(initial={'user': request.user.pk})
    }
    return direct_to_template(request, 'networks/host_form.html', extra_context)

@login_required
def host_update(request, object_id):

    host, edit = get_object_or_forbidden(Host, object_id, request.user)

    if not edit:
        raise Http404()

    return update_object(request, object_id=object_id,
                         form_class=HostUpdateForm,
                         template_name='networks/host_update.html')

@login_required
def host_delete(request, object_id):

    host, edit = get_object_or_forbidden(Host, object_id, request.user)

    if host.user != request.user:
        raise Http404()

    return delete_object(request, object_id=object_id, model=Host,
                         post_delete_redirect=reverse('host_list'))

@login_required
def network_list(request, page=None):
    search_phrase = request.GET.get('s')
    if search_phrase:
        nets = search(Network, search_phrase)
    else:
        nets = filter_user_objects(request.user, Network)

    paginator = Paginator(list(nets), 10)

    page = page or request.GET.get('page', 1)
    try:
        nets = paginator.page(page)
    except PageNotAnInteger:
        nets = paginator.page(1)
    except EmptyPage:
        nets = paginator.page(paginator.num_pages)

    extra_context = {
        'networks': nets,
        'url': '/network/network/list/'
    }
    return direct_to_template(request, 'networks/network_list.html',
                              extra_context=extra_context)

@login_required
def network_detail(request, object_id):
    """
    Network detail page has the following features:
        * displaying basic network info (name, description, etc.)
        * listing hosts related to network
        * creating relations between network and host
        * removing relations between network and host(s)
    """
    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    # remove relation between the network and selected host(s)
    if request.POST.getlist('remove_host'):
        if edit:
            hosts_pk = request.POST.getlist('remove_host')
            network_host = NetworkHost.objects.filter(network=network,
                                                      host__pk__in=hosts_pk)
            network_host.delete()

    # create relation between the network and selected host
    if request.POST.get('add_host'):
        if edit:
            host = Host.objects.get(pk=request.POST.get('add_host'))
            network_host = NetworkHost(network=network, host=host)
            network_host.save()

    queryset = Network.objects.filter(user=request.user)
    if network.hosts():
        hosts_ids = [host.pk for host in network.hosts()]
        hosts_other = Host.objects.exclude(pk__in=hosts_ids).filter(user=request.user)
    else:
        hosts_other = Host.objects.filter(user=request.user)
    extra_context = {
        'hosts_other': hosts_other,
        'can_edit': edit
    }
    return object_detail(request, queryset, object_id,
                         extra_context=extra_context)

@login_required
def network_create(request):

    if request.method == 'POST':
        form = NetworkCreateForm(request.POST)
        if form.is_valid():
            network = form.save()
            return redirect_to(request, url=network.get_absolute_url())

    extra_context = {
        'form': NetworkCreateForm(initial={'user': request.user.pk})
    }
    return direct_to_template(request, 'networks/network_form.html',
                              extra_context)

@login_required
def network_update(request, object_id):

    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    if not edit:
        raise Http404()

    return update_object(request, object_id=object_id,
                         form_class=NetworkUpdateForm,
                         template_name='networks/network_update.html')

@login_required
def network_delete(request, object_id):

    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    if network.user != request.user:
        raise Http404()

    return delete_object(request, object_id=object_id, model=Network,
                         post_delete_redirect=reverse('network_list'))

@login_required
def network_events(request, object_id):
    """Display events related to network"""

    network, edit = get_object_or_forbidden(Network, object_id, request.user)

    queryset = Network.objects.all()
    related_hosts = [nh.host.pk for nh in NetworkHost.objects.filter(network=network)]
    events = Event.objects.filter(source_host__pk__in=related_hosts)
    extra_context = {
        'events': events,
        'can_edit': edit
    }
    return object_detail(request, queryset, object_id,
                         extra_context=extra_context,
                         template_name='networks/network_events.html')

@login_required
def share(request, object_type, object_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    user_id = request.POST.get('share')
    if user_id:
        user = User.objects.get(pk=user_id)
        grant_access(obj, user)
    return share_list(request, object_type, object_id)

@login_required
def share_not(request, object_type, object_id, user_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    user = User.objects.get(pk=user_id)
    revoke_access(obj, user)
    return share_list(request, object_type, object_id)

@login_required
def share_edit(request, object_type, object_id, user_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    user = User.objects.get(pk=user_id)
    if edit:
        revoke_edit(obj, user)
    else:
        grant_edit(obj, user)
    return share_list(request, object_type, object_id)

@login_required
def share_list(request, object_type, object_id):
    model = Network if object_type == 'network' else Host
    obj, edit = get_object_or_forbidden(model, object_id, request.user)
    all_users = User.objects.exclude(pk=request.user.pk)
    other_users = []
    for user in all_users:
        if not user_has_access(obj, user):
            other_users.append(user)
    extra_context = {
        'object': obj,
        'object_type': object_type,
        'other_users': other_users
    }
    return direct_to_template(request, 'networks/share.html', extra_context)
<|endoftext|>


<commit_before>import os
import re
import sys
import warnings
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
from seleniumbase.config import settings
from seleniumbase.config import proxy_list
from seleniumbase.core import download_helper
from seleniumbase.fixtures import constants
from seleniumbase.fixtures import page_utils
import drivers  # webdriver storage folder for SeleniumBase
DRIVER_DIR = os.path.dirname(os.path.realpath(drivers.__file__))
PLATFORM = sys.platform
LOCAL_CHROMEDRIVER = None
LOCAL_GECKODRIVER = None
LOCAL_EDGEDRIVER = None
LOCAL_OPERADRIVER = None
if "darwin" in PLATFORM or "linux" in PLATFORM:
    LOCAL_CHROMEDRIVER = DRIVER_DIR + '/chromedriver'
    LOCAL_GECKODRIVER = DRIVER_DIR + '/geckodriver'
    LOCAL_OPERADRIVER = DRIVER_DIR + '/operadriver'
elif "win32" in PLATFORM or "win64" in PLATFORM or "x64" in PLATFORM:
    LOCAL_EDGEDRIVER = DRIVER_DIR + '/MicrosoftWebDriver.exe'
    LOCAL_CHROMEDRIVER = DRIVER_DIR + '/chromedriver.exe'
    LOCAL_GECKODRIVER = DRIVER_DIR + '/geckodriver.exe'
    LOCAL_OPERADRIVER = DRIVER_DIR + '/operadriver.exe'
else:
    # Cannot determine system
    pass  # SeleniumBase will use web drivers from the System PATH by default

def make_executable(file_path):
    # Set permissions to: "If you can read it, you can execute it."
    mode = os.stat(file_path).st_mode
    mode |= (mode & 0o444) >> 2  # copy R bits to X
    os.chmod(file_path, mode)

def make_driver_executable_if_not(driver_path):
    # Verify driver has executable permissions. If not, add them.
    permissions = oct(os.stat(driver_path)[0])[-3:]
    if '4' in permissions or '6' in permissions:
        # We want at least a '5' or '7' to make sure it's executable
        make_executable(driver_path)

def _set_chrome_options(downloads_path, proxy_string):
    chrome_options = webdriver.ChromeOptions()
    prefs = {
        "download.default_directory": downloads_path,
        "credentials_enable_service": False,
        "profile": {
            "password_manager_enabled": False
        }
    }
    chrome_options.add_experimental_option("prefs", prefs)
    chrome_options.add_argument("--test-type")
    chrome_options.add_argument("--no-first-run")
    chrome_options.add_argument("--ignore-certificate-errors")
    chrome_options.add_argument("--allow-file-access-from-files")
    chrome_options.add_argument("--allow-insecure-localhost")
    chrome_options.add_argument("--allow-running-insecure-content")
    chrome_options.add_argument("--disable-infobars")
    chrome_options.add_argument("--disable-save-password-bubble")
    chrome_options.add_argument("--disable-single-click-autofill")
    chrome_options.add_argument("--disable-translate")
    chrome_options.add_argument("--disable-web-security")
    if proxy_string:
        chrome_options.add_argument('--proxy-server=%s' % proxy_string)
    if settings.START_CHROME_IN_FULL_SCREEN_MODE:
        # Run Chrome in full screen mode on WINDOWS
        chrome_options.add_argument("--start-maximized")
        # Run Chrome in full screen mode on MAC/Linux
        chrome_options.add_argument("--kiosk")
    if "win32" in sys.platform or "win64" in sys.platform:
        chrome_options.add_argument("--log-level=3")
    return chrome_options

def _create_firefox_profile(downloads_path, proxy_string):
    profile = webdriver.FirefoxProfile()
    profile.set_preference("reader.parse-on-load.enabled", False)
    profile.set_preference("pdfjs.disabled", True)
    profile.set_preference("app.update.auto", False)
    profile.set_preference("app.update.enabled", False)
    if proxy_string:
        proxy_server = proxy_string.split(':')[0]
        proxy_port = proxy_string.split(':')[1]
        profile.set_preference("network.proxy.type", 1)
        profile.set_preference("network.proxy.http", proxy_server)
        profile.set_preference("network.proxy.http_port", int(proxy_port))
        profile.set_preference("network.proxy.ssl", proxy_server)
        profile.set_preference("network.proxy.ssl_port", int(proxy_port))
    profile.set_preference(
        "security.mixed_content.block_active_content", False)
    profile.set_preference("security.csp.enable", False)
    profile.set_preference(
        "browser.download.manager.showAlertOnComplete", False)
    profile.set_preference("browser.download.panel.shown", False)
    profile.set_preference(
        "browser.download.animateNotifications", False)
    profile.set_preference("browser.download.dir", downloads_path)
    profile.set_preference("browser.download.folderList", 2)
    profile.set_preference(
        "browser.helperApps.neverAsk.saveToDisk",
        ("application/pdf, application/zip, application/octet-stream, "
         "text/csv, text/xml, application/xml, text/plain, "
         "text/octet-stream, "
         "application/"
         "vnd.openxmlformats-officedocument.spreadsheetml.sheet"))
    return profile

def display_proxy_warning(proxy_string):
    message = ('\n\nWARNING: Proxy String ["%s"] is NOT in the expected '
               '"ip_address:port" or "server:port" format, '
               '(OR the key does not exist in proxy_list.PROXY_LIST). '
               '*** DEFAULTING to NOT USING a Proxy Server! ***'
               % proxy_string)
    warnings.simplefilter('always', Warning)  # See Warnings
    warnings.warn(message, category=Warning, stacklevel=2)
    warnings.simplefilter('default', Warning)  # Set Default

def validate_proxy_string(proxy_string):
    if proxy_string in proxy_list.PROXY_LIST.keys():
        proxy_string = proxy_list.PROXY_LIST[proxy_string]
        if not proxy_string:
            return None
    valid = False
    val_ip = re.match('^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}:\d+$', proxy_string)
    if not val_ip:
        if proxy_string.startswith('http://'):
            proxy_string = proxy_string.split('http://')[1]
        elif proxy_string.startswith('https://'):
            proxy_string = proxy_string.split('https://')[1]
        elif '://' in proxy_string:
            proxy_string = proxy_string.split('://')[1]
        chunks = proxy_string.split(':')
        if len(chunks) == 2:
            if re.match('^\d+$', chunks[1]):
                if page_utils.is_valid_url('http://' + proxy_string):
                    valid = True
    else:
        proxy_string = val_ip.group()
        valid = True
    if not valid:
        display_proxy_warning(proxy_string)
        proxy_string = None
    return proxy_string

def get_driver(browser_name, headless=False, use_grid=False,
               servername='localhost', port=4444, proxy_string=None):
    if proxy_string:
        proxy_string = validate_proxy_string(proxy_string)
    if use_grid:
        return get_remote_driver(
            browser_name, headless, servername, port, proxy_string)
    else:
        return get_local_driver(browser_name, headless, proxy_string)

def get_remote_driver(browser_name, headless, servername, port, proxy_string):
    downloads_path = download_helper.get_downloads_folder()
    download_helper.reset_downloads_folder()
    address = "http://%s:%s/wd/hub" % (servername, port)

    if browser_name == constants.Browser.GOOGLE_CHROME:
        chrome_options = _set_chrome_options(downloads_path, proxy_string)
        if headless:
            chrome_options.add_argument("--headless")
            chrome_options.add_argument("--disable-gpu")
            chrome_options.add_argument("--no-sandbox")
        capabilities = chrome_options.to_capabilities()
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=capabilities)
    elif browser_name == constants.Browser.FIREFOX:
        try:
            # Use Geckodriver for Firefox if it's on the PATH
            profile = _create_firefox_profile(downloads_path, proxy_string)
            firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
            firefox_capabilities['marionette'] = True
            if headless:
                firefox_capabilities['moz:firefoxOptions'] = (
                    {'args': ['-headless']})
            capabilities = firefox_capabilities
            address = "http://%s:%s/wd/hub" % (servername, port)
            return webdriver.Remote(
                command_executor=address,
                desired_capabilities=capabilities,
                browser_profile=profile)
        except WebDriverException:
            # Don't use Geckodriver: Only works for old versions of Firefox
            profile = _create_firefox_profile(downloads_path, proxy_string)
            firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
            firefox_capabilities['marionette'] = False
            if headless:
                firefox_capabilities['moz:firefoxOptions'] = (
                    {'args': ['-headless']})
            capabilities = firefox_capabilities
            return webdriver.Remote(
                command_executor=address,
                desired_capabilities=capabilities,
                browser_profile=profile)
    elif browser_name == constants.Browser.INTERNET_EXPLORER:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.INTERNETEXPLORER))
    elif browser_name == constants.Browser.EDGE:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.EDGE))
    elif browser_name == constants.Browser.SAFARI:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.SAFARI))
    elif browser_name == constants.Browser.OPERA:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.OPERA))
    elif browser_name == constants.Browser.PHANTOM_JS:
        with warnings.catch_warnings():
            # Ignore "PhantomJS has been deprecated" UserWarning
            warnings.simplefilter("ignore", category=UserWarning)
            return webdriver.Remote(
                command_executor=address,
                desired_capabilities=(
                    webdriver.DesiredCapabilities.PHANTOMJS))

def get_local_driver(browser_name, headless, proxy_string):
    '''
    Spins up a new web browser and returns the driver.
    Can also be used to spin up additional browsers for the same test.
    '''
    downloads_path = download_helper.get_downloads_folder()
    download_helper.reset_downloads_folder()

    if browser_name == constants.Browser.FIREFOX:
        try:
            try:
                # Use Geckodriver for Firefox if it's on the PATH
                profile = _create_firefox_profile(downloads_path, proxy_string)
                firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                firefox_capabilities['marionette'] = True
                options = webdriver.FirefoxOptions()
                if headless:
                    options.add_argument('-headless')
                if LOCAL_GECKODRIVER and os.path.exists(LOCAL_GECKODRIVER):
                    make_driver_executable_if_not(LOCAL_GECKODRIVER)
                    firefox_driver = webdriver.Firefox(
                        firefox_profile=profile,
                        capabilities=firefox_capabilities,
                        firefox_options=options,
                        executable_path=LOCAL_GECKODRIVER)
                else:
                    firefox_driver = webdriver.Firefox(
                        firefox_profile=profile,
                        capabilities=firefox_capabilities,
                        firefox_options=options)
            except WebDriverException:
                # Don't use Geckodriver: Only works for old versions of Firefox
                profile = _create_firefox_profile(downloads_path, proxy_string)
                firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                firefox_capabilities['marionette'] = False
                firefox_driver = webdriver.Firefox(
                    firefox_profile=profile, capabilities=firefox_capabilities)
            return firefox_driver
        except Exception as e:
            if headless:
                raise Exception(e)
            return webdriver.Firefox()
    elif browser_name == constants.Browser.INTERNET_EXPLORER:
        return webdriver.Ie()
    elif browser_name == constants.Browser.EDGE:
        edge_capabilities = DesiredCapabilities.EDGE.copy()
        if LOCAL_EDGEDRIVER and os.path.exists(LOCAL_EDGEDRIVER):
            make_driver_executable_if_not(LOCAL_EDGEDRIVER)
            return webdriver.Edge(
                capabilities=edge_capabilities,
                executable_path=LOCAL_EDGEDRIVER)
        else:
            return webdriver.Edge(capabilities=edge_capabilities)
    elif browser_name == constants.Browser.SAFARI:
        return webdriver.Safari()
    elif browser_name == constants.Browser.OPERA:
        if LOCAL_OPERADRIVER and os.path.exists(LOCAL_OPERADRIVER):
            make_driver_executable_if_not(LOCAL_OPERADRIVER)
            return webdriver.Opera(executable_path=LOCAL_OPERADRIVER)
        else:
            return webdriver.Opera()
    elif browser_name == constants.Browser.PHANTOM_JS:
        with warnings.catch_warnings():
            # Ignore "PhantomJS has been deprecated" UserWarning
            warnings.simplefilter("ignore", category=UserWarning)
            return webdriver.PhantomJS()
    elif browser_name == constants.Browser.GOOGLE_CHROME:
        try:
            chrome_options = _set_chrome_options(downloads_path, proxy_string)
            if headless:
                chrome_options.add_argument("--headless")
                chrome_options.add_argument("--disable-gpu")
                chrome_options.add_argument("--no-sandbox")
            if LOCAL_CHROMEDRIVER and os.path.exists(LOCAL_CHROMEDRIVER):
                make_driver_executable_if_not(LOCAL_CHROMEDRIVER)
                return webdriver.Chrome(
                    executable_path=LOCAL_CHROMEDRIVER, options=chrome_options)
            else:
                return webdriver.Chrome(options=chrome_options)
        except Exception as e:
            if headless:
                raise Exception(e)
            if LOCAL_CHROMEDRIVER and os.path.exists(LOCAL_CHROMEDRIVER):
                make_driver_executable_if_not(LOCAL_CHROMEDRIVER)
                return webdriver.Chrome(executable_path=LOCAL_CHROMEDRIVER)
            else:
                return webdriver.Chrome()
<commit_msg>Use the new "drivers" folder to finding web drivers<commit_after>import os
import re
import sys
import warnings
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
from seleniumbase.config import settings
from seleniumbase.config import proxy_list
from seleniumbase.core import download_helper
from seleniumbase.fixtures import constants
from seleniumbase.fixtures import page_utils
from seleniumbase import drivers  # webdriver storage folder for SeleniumBase
DRIVER_DIR = os.path.dirname(os.path.realpath(drivers.__file__))
PLATFORM = sys.platform
LOCAL_CHROMEDRIVER = None
LOCAL_GECKODRIVER = None
LOCAL_EDGEDRIVER = None
LOCAL_OPERADRIVER = None
if "darwin" in PLATFORM or "linux" in PLATFORM:
    LOCAL_CHROMEDRIVER = DRIVER_DIR + '/chromedriver'
    LOCAL_GECKODRIVER = DRIVER_DIR + '/geckodriver'
    LOCAL_OPERADRIVER = DRIVER_DIR + '/operadriver'
elif "win32" in PLATFORM or "win64" in PLATFORM or "x64" in PLATFORM:
    LOCAL_EDGEDRIVER = DRIVER_DIR + '/MicrosoftWebDriver.exe'
    LOCAL_CHROMEDRIVER = DRIVER_DIR + '/chromedriver.exe'
    LOCAL_GECKODRIVER = DRIVER_DIR + '/geckodriver.exe'
    LOCAL_OPERADRIVER = DRIVER_DIR + '/operadriver.exe'
else:
    # Cannot determine system
    pass  # SeleniumBase will use web drivers from the System PATH by default

def make_executable(file_path):
    # Set permissions to: "If you can read it, you can execute it."
    mode = os.stat(file_path).st_mode
    mode |= (mode & 0o444) >> 2  # copy R bits to X
    os.chmod(file_path, mode)

def make_driver_executable_if_not(driver_path):
    # Verify driver has executable permissions. If not, add them.
    permissions = oct(os.stat(driver_path)[0])[-3:]
    if '4' in permissions or '6' in permissions:
        # We want at least a '5' or '7' to make sure it's executable
        make_executable(driver_path)

def _set_chrome_options(downloads_path, proxy_string):
    chrome_options = webdriver.ChromeOptions()
    prefs = {
        "download.default_directory": downloads_path,
        "credentials_enable_service": False,
        "profile": {
            "password_manager_enabled": False
        }
    }
    chrome_options.add_experimental_option("prefs", prefs)
    chrome_options.add_argument("--test-type")
    chrome_options.add_argument("--no-first-run")
    chrome_options.add_argument("--ignore-certificate-errors")
    chrome_options.add_argument("--allow-file-access-from-files")
    chrome_options.add_argument("--allow-insecure-localhost")
    chrome_options.add_argument("--allow-running-insecure-content")
    chrome_options.add_argument("--disable-infobars")
    chrome_options.add_argument("--disable-save-password-bubble")
    chrome_options.add_argument("--disable-single-click-autofill")
    chrome_options.add_argument("--disable-translate")
    chrome_options.add_argument("--disable-web-security")
    if proxy_string:
        chrome_options.add_argument('--proxy-server=%s' % proxy_string)
    if settings.START_CHROME_IN_FULL_SCREEN_MODE:
        # Run Chrome in full screen mode on WINDOWS
        chrome_options.add_argument("--start-maximized")
        # Run Chrome in full screen mode on MAC/Linux
        chrome_options.add_argument("--kiosk")
    if "win32" in sys.platform or "win64" in sys.platform:
        chrome_options.add_argument("--log-level=3")
    return chrome_options

def _create_firefox_profile(downloads_path, proxy_string):
    profile = webdriver.FirefoxProfile()
    profile.set_preference("reader.parse-on-load.enabled", False)
    profile.set_preference("pdfjs.disabled", True)
    profile.set_preference("app.update.auto", False)
    profile.set_preference("app.update.enabled", False)
    if proxy_string:
        proxy_server = proxy_string.split(':')[0]
        proxy_port = proxy_string.split(':')[1]
        profile.set_preference("network.proxy.type", 1)
        profile.set_preference("network.proxy.http", proxy_server)
        profile.set_preference("network.proxy.http_port", int(proxy_port))
        profile.set_preference("network.proxy.ssl", proxy_server)
        profile.set_preference("network.proxy.ssl_port", int(proxy_port))
    profile.set_preference(
        "security.mixed_content.block_active_content", False)
    profile.set_preference("security.csp.enable", False)
    profile.set_preference(
        "browser.download.manager.showAlertOnComplete", False)
    profile.set_preference("browser.download.panel.shown", False)
    profile.set_preference(
        "browser.download.animateNotifications", False)
    profile.set_preference("browser.download.dir", downloads_path)
    profile.set_preference("browser.download.folderList", 2)
    profile.set_preference(
        "browser.helperApps.neverAsk.saveToDisk",
        ("application/pdf, application/zip, application/octet-stream, "
         "text/csv, text/xml, application/xml, text/plain, "
         "text/octet-stream, "
         "application/"
         "vnd.openxmlformats-officedocument.spreadsheetml.sheet"))
    return profile

def display_proxy_warning(proxy_string):
    message = ('\n\nWARNING: Proxy String ["%s"] is NOT in the expected '
               '"ip_address:port" or "server:port" format, '
               '(OR the key does not exist in proxy_list.PROXY_LIST). '
               '*** DEFAULTING to NOT USING a Proxy Server! ***'
               % proxy_string)
    warnings.simplefilter('always', Warning)  # See Warnings
    warnings.warn(message, category=Warning, stacklevel=2)
    warnings.simplefilter('default', Warning)  # Set Default

def validate_proxy_string(proxy_string):
    if proxy_string in proxy_list.PROXY_LIST.keys():
        proxy_string = proxy_list.PROXY_LIST[proxy_string]
        if not proxy_string:
            return None
    valid = False
    val_ip = re.match('^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}:\d+$', proxy_string)
    if not val_ip:
        if proxy_string.startswith('http://'):
            proxy_string = proxy_string.split('http://')[1]
        elif proxy_string.startswith('https://'):
            proxy_string = proxy_string.split('https://')[1]
        elif '://' in proxy_string:
            proxy_string = proxy_string.split('://')[1]
        chunks = proxy_string.split(':')
        if len(chunks) == 2:
            if re.match('^\d+$', chunks[1]):
                if page_utils.is_valid_url('http://' + proxy_string):
                    valid = True
    else:
        proxy_string = val_ip.group()
        valid = True
    if not valid:
        display_proxy_warning(proxy_string)
        proxy_string = None
    return proxy_string

def get_driver(browser_name, headless=False, use_grid=False,
               servername='localhost', port=4444, proxy_string=None):
    if proxy_string:
        proxy_string = validate_proxy_string(proxy_string)
    if use_grid:
        return get_remote_driver(
            browser_name, headless, servername, port, proxy_string)
    else:
        return get_local_driver(browser_name, headless, proxy_string)

def get_remote_driver(browser_name, headless, servername, port, proxy_string):
    downloads_path = download_helper.get_downloads_folder()
    download_helper.reset_downloads_folder()
    address = "http://%s:%s/wd/hub" % (servername, port)

    if browser_name == constants.Browser.GOOGLE_CHROME:
        chrome_options = _set_chrome_options(downloads_path, proxy_string)
        if headless:
            chrome_options.add_argument("--headless")
            chrome_options.add_argument("--disable-gpu")
            chrome_options.add_argument("--no-sandbox")
        capabilities = chrome_options.to_capabilities()
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=capabilities)
    elif browser_name == constants.Browser.FIREFOX:
        try:
            # Use Geckodriver for Firefox if it's on the PATH
            profile = _create_firefox_profile(downloads_path, proxy_string)
            firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
            firefox_capabilities['marionette'] = True
            if headless:
                firefox_capabilities['moz:firefoxOptions'] = (
                    {'args': ['-headless']})
            capabilities = firefox_capabilities
            address = "http://%s:%s/wd/hub" % (servername, port)
            return webdriver.Remote(
                command_executor=address,
                desired_capabilities=capabilities,
                browser_profile=profile)
        except WebDriverException:
            # Don't use Geckodriver: Only works for old versions of Firefox
            profile = _create_firefox_profile(downloads_path, proxy_string)
            firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
            firefox_capabilities['marionette'] = False
            if headless:
                firefox_capabilities['moz:firefoxOptions'] = (
                    {'args': ['-headless']})
            capabilities = firefox_capabilities
            return webdriver.Remote(
                command_executor=address,
                desired_capabilities=capabilities,
                browser_profile=profile)
    elif browser_name == constants.Browser.INTERNET_EXPLORER:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.INTERNETEXPLORER))
    elif browser_name == constants.Browser.EDGE:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.EDGE))
    elif browser_name == constants.Browser.SAFARI:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.SAFARI))
    elif browser_name == constants.Browser.OPERA:
        return webdriver.Remote(
            command_executor=address,
            desired_capabilities=(
                webdriver.DesiredCapabilities.OPERA))
    elif browser_name == constants.Browser.PHANTOM_JS:
        with warnings.catch_warnings():
            # Ignore "PhantomJS has been deprecated" UserWarning
            warnings.simplefilter("ignore", category=UserWarning)
            return webdriver.Remote(
                command_executor=address,
                desired_capabilities=(
                    webdriver.DesiredCapabilities.PHANTOMJS))

def get_local_driver(browser_name, headless, proxy_string):
    '''
    Spins up a new web browser and returns the driver.
    Can also be used to spin up additional browsers for the same test.
    '''
    downloads_path = download_helper.get_downloads_folder()
    download_helper.reset_downloads_folder()

    if browser_name == constants.Browser.FIREFOX:
        try:
            try:
                # Use Geckodriver for Firefox if it's on the PATH
                profile = _create_firefox_profile(downloads_path, proxy_string)
                firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                firefox_capabilities['marionette'] = True
                options = webdriver.FirefoxOptions()
                if headless:
                    options.add_argument('-headless')
                if LOCAL_GECKODRIVER and os.path.exists(LOCAL_GECKODRIVER):
                    make_driver_executable_if_not(LOCAL_GECKODRIVER)
                    firefox_driver = webdriver.Firefox(
                        firefox_profile=profile,
                        capabilities=firefox_capabilities,
                        firefox_options=options,
                        executable_path=LOCAL_GECKODRIVER)
                else:
                    firefox_driver = webdriver.Firefox(
                        firefox_profile=profile,
                        capabilities=firefox_capabilities,
                        firefox_options=options)
            except WebDriverException:
                # Don't use Geckodriver: Only works for old versions of Firefox
                profile = _create_firefox_profile(downloads_path, proxy_string)
                firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                firefox_capabilities['marionette'] = False
                firefox_driver = webdriver.Firefox(
                    firefox_profile=profile, capabilities=firefox_capabilities)
            return firefox_driver
        except Exception as e:
            if headless:
                raise Exception(e)
            return webdriver.Firefox()
    elif browser_name == constants.Browser.INTERNET_EXPLORER:
        return webdriver.Ie()
    elif browser_name == constants.Browser.EDGE:
        edge_capabilities = DesiredCapabilities.EDGE.copy()
        if LOCAL_EDGEDRIVER and os.path.exists(LOCAL_EDGEDRIVER):
            make_driver_executable_if_not(LOCAL_EDGEDRIVER)
            return webdriver.Edge(
                capabilities=edge_capabilities,
                executable_path=LOCAL_EDGEDRIVER)
        else:
            return webdriver.Edge(capabilities=edge_capabilities)
    elif browser_name == constants.Browser.SAFARI:
        return webdriver.Safari()
    elif browser_name == constants.Browser.OPERA:
        if LOCAL_OPERADRIVER and os.path.exists(LOCAL_OPERADRIVER):
            make_driver_executable_if_not(LOCAL_OPERADRIVER)
            return webdriver.Opera(executable_path=LOCAL_OPERADRIVER)
        else:
            return webdriver.Opera()
    elif browser_name == constants.Browser.PHANTOM_JS:
        with warnings.catch_warnings():
            # Ignore "PhantomJS has been deprecated" UserWarning
            warnings.simplefilter("ignore", category=UserWarning)
            return webdriver.PhantomJS()
    elif browser_name == constants.Browser.GOOGLE_CHROME:
        try:
            chrome_options = _set_chrome_options(downloads_path, proxy_string)
            if headless:
                chrome_options.add_argument("--headless")
                chrome_options.add_argument("--disable-gpu")
                chrome_options.add_argument("--no-sandbox")
            if LOCAL_CHROMEDRIVER and os.path.exists(LOCAL_CHROMEDRIVER):
                make_driver_executable_if_not(LOCAL_CHROMEDRIVER)
                return webdriver.Chrome(
                    executable_path=LOCAL_CHROMEDRIVER, options=chrome_options)
            else:
                return webdriver.Chrome(options=chrome_options)
        except Exception as e:
            if headless:
                raise Exception(e)
            if LOCAL_CHROMEDRIVER and os.path.exists(LOCAL_CHROMEDRIVER):
                make_driver_executable_if_not(LOCAL_CHROMEDRIVER)
                return webdriver.Chrome(executable_path=LOCAL_CHROMEDRIVER)
            else:
                return webdriver.Chrome()
<|endoftext|>


<commit_before>import time
from collections import defaultdict
from functools import partial
from typing import Optional

import cereal.messaging as messaging
from system.swaglog import cloudlog
from selfdrive.boardd.boardd import can_list_to_can_capnp
from panda.python.uds import CanClient, IsoTpMessage, FUNCTIONAL_ADDRS, get_rx_addr_for_tx_addr

class IsoTpParallelQuery:
  def __init__(self, sendcan, logcan, bus, addrs, request, response, response_offset=0x8, functional_addr=False, debug=False, response_pending_timeout=10):
    self.sendcan = sendcan
    self.logcan = logcan
    self.bus = bus
    self.request = request
    self.response = response
    self.debug = debug
    self.functional_addr = functional_addr
    self.response_pending_timeout = response_pending_timeout

    self.real_addrs = []
    for a in addrs:
      if isinstance(a, tuple):
        self.real_addrs.append(a)
      else:
        self.real_addrs.append((a, None))

    self.msg_addrs = {tx_addr: get_rx_addr_for_tx_addr(tx_addr[0], rx_offset=response_offset) for tx_addr in self.real_addrs}
    self.msg_buffer = defaultdict(list)

  def rx(self):
    """Drain can socket and sort messages into buffers based on address"""
    can_packets = messaging.drain_sock(self.logcan, wait_for_one=True)

    for packet in can_packets:
      for msg in packet.can:
        if msg.src == self.bus:
          if self.functional_addr:
            if (0x7E8 <= msg.address <= 0x7EF) or (0x18DAF100 <= msg.address <= 0x18DAF1FF):
              fn_addr = next(a for a in FUNCTIONAL_ADDRS if msg.address - a <= 32)
              self.msg_buffer[fn_addr].append((msg.address, msg.busTime, msg.dat, msg.src))
          elif msg.address in self.msg_addrs.values():
            self.msg_buffer[msg.address].append((msg.address, msg.busTime, msg.dat, msg.src))

  def _can_tx(self, tx_addr, dat, bus):
    """Helper function to send single message"""
    msg = [tx_addr, 0, dat, bus]
    self.sendcan.send(can_list_to_can_capnp([msg], msgtype='sendcan'))

  def _can_rx(self, addr, sub_addr=None):
    """Helper function to retrieve message with specified address and subadress from buffer"""
    keep_msgs = []

    if sub_addr is None:
      msgs = self.msg_buffer[addr]
    else:
      # Filter based on subadress
      msgs = []
      for m in self.msg_buffer[addr]:
        first_byte = m[2][0]
        if first_byte == sub_addr:
          msgs.append(m)
        else:
          keep_msgs.append(m)

    self.msg_buffer[addr] = keep_msgs
    return msgs

  def _drain_rx(self):
    messaging.drain_sock(self.logcan)
    self.msg_buffer = defaultdict(list)

  def get_data(self, timeout, total_timeout=60.):
    self._drain_rx()

    # Create message objects
    msgs = {}
    request_counter = {}
    request_done = {}
    for tx_addr, rx_addr in self.msg_addrs.items():
      # rx_addr not set when using functional tx addr
      id_addr = rx_addr or tx_addr[0]
      sub_addr = tx_addr[1]

      can_client = CanClient(self._can_tx, partial(self._can_rx, id_addr, sub_addr=sub_addr), tx_addr[0], rx_addr,
                             self.bus, sub_addr=sub_addr, debug=self.debug)

      max_len = 8 if sub_addr is None else 7

      msg = IsoTpMessage(can_client, timeout=0, max_len=max_len, debug=self.debug)
      msg.send(self.request[0])

      msgs[tx_addr] = msg
      request_counter[tx_addr] = 0
      request_done[tx_addr] = False

    results = {}
    start_time = time.monotonic()
    response_timeouts = {tx_addr: start_time + timeout for tx_addr in self.msg_addrs}
    while True:
      self.rx()

      if all(request_done.values()):
        break

      for tx_addr, msg in msgs.items():
        try:
          dat: Optional[bytes] = msg.recv()
        except Exception:
          cloudlog.exception("Error processing UDS response")
          request_done[tx_addr] = True
          continue

        if not dat:
          continue

        counter = request_counter[tx_addr]
        expected_response = self.response[counter]
        response_valid = dat[:len(expected_response)] == expected_response

        if response_valid:
          response_timeouts[tx_addr] = time.monotonic() + timeout
          if counter + 1 < len(self.request):
            msg.send(self.request[counter + 1])
            request_counter[tx_addr] += 1
          else:
            results[(tx_addr, msg._can_client.rx_addr)] = dat[len(expected_response):]
            request_done[tx_addr] = True
        else:
          error_code = dat[2] if len(dat) > 2 else -1
          if error_code == 0x78:
            response_timeouts[tx_addr] = time.monotonic() + self.response_pending_timeout
            if self.debug:
              cloudlog.warning(f"iso-tp query response pending: {tx_addr}")
          else:
            request_done[tx_addr] = True
            cloudlog.warning(f"iso-tp query bad response: {tx_addr} - 0x{dat.hex()}")

      cur_time = time.monotonic()
      if cur_time - max(response_timeouts.values()) > 0:
        for tx_addr in msgs:
          if request_counter[tx_addr] > 0 and not request_done[tx_addr]:
            cloudlog.warning(f"iso-tp query timeout after receiving response: {tx_addr}")
        break

      if cur_time - start_time > total_timeout:
        cloudlog.warning("iso-tp query timeout while receiving data")
        break

    return results
<commit_msg>Clean up timers on receipt of a non-0x78 UDS error (#25398)<commit_after>import time
from collections import defaultdict
from functools import partial
from typing import Optional

import cereal.messaging as messaging
from system.swaglog import cloudlog
from selfdrive.boardd.boardd import can_list_to_can_capnp
from panda.python.uds import CanClient, IsoTpMessage, FUNCTIONAL_ADDRS, get_rx_addr_for_tx_addr

class IsoTpParallelQuery:
  def __init__(self, sendcan, logcan, bus, addrs, request, response, response_offset=0x8, functional_addr=False, debug=False, response_pending_timeout=10):
    self.sendcan = sendcan
    self.logcan = logcan
    self.bus = bus
    self.request = request
    self.response = response
    self.debug = debug
    self.functional_addr = functional_addr
    self.response_pending_timeout = response_pending_timeout

    self.real_addrs = []
    for a in addrs:
      if isinstance(a, tuple):
        self.real_addrs.append(a)
      else:
        self.real_addrs.append((a, None))

    self.msg_addrs = {tx_addr: get_rx_addr_for_tx_addr(tx_addr[0], rx_offset=response_offset) for tx_addr in self.real_addrs}
    self.msg_buffer = defaultdict(list)

  def rx(self):
    """Drain can socket and sort messages into buffers based on address"""
    can_packets = messaging.drain_sock(self.logcan, wait_for_one=True)

    for packet in can_packets:
      for msg in packet.can:
        if msg.src == self.bus:
          if self.functional_addr:
            if (0x7E8 <= msg.address <= 0x7EF) or (0x18DAF100 <= msg.address <= 0x18DAF1FF):
              fn_addr = next(a for a in FUNCTIONAL_ADDRS if msg.address - a <= 32)
              self.msg_buffer[fn_addr].append((msg.address, msg.busTime, msg.dat, msg.src))
          elif msg.address in self.msg_addrs.values():
            self.msg_buffer[msg.address].append((msg.address, msg.busTime, msg.dat, msg.src))

  def _can_tx(self, tx_addr, dat, bus):
    """Helper function to send single message"""
    msg = [tx_addr, 0, dat, bus]
    self.sendcan.send(can_list_to_can_capnp([msg], msgtype='sendcan'))

  def _can_rx(self, addr, sub_addr=None):
    """Helper function to retrieve message with specified address and subadress from buffer"""
    keep_msgs = []

    if sub_addr is None:
      msgs = self.msg_buffer[addr]
    else:
      # Filter based on subadress
      msgs = []
      for m in self.msg_buffer[addr]:
        first_byte = m[2][0]
        if first_byte == sub_addr:
          msgs.append(m)
        else:
          keep_msgs.append(m)

    self.msg_buffer[addr] = keep_msgs
    return msgs

  def _drain_rx(self):
    messaging.drain_sock(self.logcan)
    self.msg_buffer = defaultdict(list)

  def get_data(self, timeout, total_timeout=60.):
    self._drain_rx()

    # Create message objects
    msgs = {}
    request_counter = {}
    request_done = {}
    for tx_addr, rx_addr in self.msg_addrs.items():
      # rx_addr not set when using functional tx addr
      id_addr = rx_addr or tx_addr[0]
      sub_addr = tx_addr[1]

      can_client = CanClient(self._can_tx, partial(self._can_rx, id_addr, sub_addr=sub_addr), tx_addr[0], rx_addr,
                             self.bus, sub_addr=sub_addr, debug=self.debug)

      max_len = 8 if sub_addr is None else 7

      msg = IsoTpMessage(can_client, timeout=0, max_len=max_len, debug=self.debug)
      msg.send(self.request[0])

      msgs[tx_addr] = msg
      request_counter[tx_addr] = 0
      request_done[tx_addr] = False

    results = {}
    start_time = time.monotonic()
    response_timeouts = {tx_addr: start_time + timeout for tx_addr in self.msg_addrs}
    while True:
      self.rx()

      if all(request_done.values()):
        break

      for tx_addr, msg in msgs.items():
        try:
          dat: Optional[bytes] = msg.recv()
        except Exception:
          cloudlog.exception("Error processing UDS response")
          request_done[tx_addr] = True
          continue

        if not dat:
          continue

        counter = request_counter[tx_addr]
        expected_response = self.response[counter]
        response_valid = dat[:len(expected_response)] == expected_response

        if response_valid:
          response_timeouts[tx_addr] = time.monotonic() + timeout
          if counter + 1 < len(self.request):
            msg.send(self.request[counter + 1])
            request_counter[tx_addr] += 1
          else:
            results[(tx_addr, msg._can_client.rx_addr)] = dat[len(expected_response):]
            request_done[tx_addr] = True
        else:
          error_code = dat[2] if len(dat) > 2 else -1
          if error_code == 0x78:
            response_timeouts[tx_addr] = time.monotonic() + self.response_pending_timeout
            if self.debug:
              cloudlog.warning(f"iso-tp query response pending: {tx_addr}")
          else:
            response_timeouts[tx_addr] = 0
            request_done[tx_addr] = True
            cloudlog.warning(f"iso-tp query bad response: {tx_addr} - 0x{dat.hex()}")

      cur_time = time.monotonic()
      if cur_time - max(response_timeouts.values()) > 0:
        for tx_addr in msgs:
          if request_counter[tx_addr] > 0 and not request_done[tx_addr]:
            cloudlog.warning(f"iso-tp query timeout after receiving response: {tx_addr}")
        break

      if cur_time - start_time > total_timeout:
        cloudlog.warning("iso-tp query timeout while receiving data")
        break

    return results
<|endoftext|>