<commit_before>from future import standard_library
standard_library.install_aliases()
from builtins import str
import datetime
import logging
import pycurl
import io
import os
import re
import urllib.request, urllib.parse, urllib.error
import hashlib

from biomaj.download.interface import DownloadInterface
from biomaj.download.ftp import FTPDownload
from biomaj.utils import Utils

try:
    from io import BytesIO
except ImportError:
    from StringIO import StringIO as BytesIO

class MultiDownload(DownloadInterface):
    '''
    Base interface for a downloader using multiple downloaders
    '''
    def __init__(self):
        DownloadInterface.__init__(self)
        self.downloaders = []
        self.files_to_download = []

    def add_downloaders(self, downloaders):
        '''
        Adds a list of downloaders
        '''
        self.downloaders += downloaders
        for d in downloaders:
            self.files_to_download += d.files_to_download

    def match(self, patterns, file_list, dir_list=None, prefix='', submatch=False):
        if dir_list is None:
            dir_list = []
        self.files_to_download = []
        for d in self.downloaders:
            d.match(patterns, d.files_to_download, [], prefix, submatch)
        self.files_to_download = []
        for d in self.downloaders:
            self.files_to_download += d.files_to_download

    def download(self, local_dir):
        self.files_to_download = []
        for d in self.downloaders:
            if self.kill_received:
                raise Exception('Kill request received, exiting')
            d.download(local_dir)
        self.files_to_download = []
        for d in self.downloaders:
            self.files_to_download += d.files_to_download
        return (self.files_to_download, [])

    def list(self):
        self.files_to_download = []
        for d in self.downloaders:
            d.list()
        self.files_to_download = []
        for d in self.downloaders:
            self.files_to_download += d.files_to_download
        return (self.files_to_download, [])

    def close(self):
        for d in self.downloaders:
            d.close()

class DirectFTPDownload(FTPDownload):
    '''
    download a list of files from FTP, no regexp
    '''

    def __init__(self, protocol, host, rootdir='', file_list=None):
        '''

        Initialize the files in list with today as last-modification date.
        Size is also preset to zero, size will be set after download

        :param file_list: list of files to download on server
        :type file_list: list
        '''
        FTPDownload.__init__(self, protocol, host, rootdir)
        if file_list is None:
            file_list = []
        today = datetime.date.today()
        self.files_to_download = []
        self.headers = {}
        for file in file_list:
            rfile = {}
            rfile['root'] = self.rootdir
            rfile['permissions'] = ''
            rfile['group'] = ''
            rfile['user'] = ''
            rfile['size'] = 0
            rfile['month'] = today.month
            rfile['day'] = today.day
            rfile['year'] = today.year
            rfile['name'] = file
            rfile['hash'] = None
            self.files_to_download.append(rfile)

    def list(self, directory=''):
        '''
        FTP protocol does not give us the possibility to get file date from remote url
        '''
        return (self.files_to_download, [])

    def match(self, patterns, file_list, dir_list=None, prefix='', submatch=False):
        '''
        All files to download match, no pattern
        '''
        if dir_list is None:
            dir_list = []
        self.files_to_download = file_list

class DirectHttpDownload(DirectFTPDownload):

    def __init__(self, protocol, host, rootdir='', file_list=None):
        '''
        :param file_list: list of files to download on server
        :type file_list: list
        '''
        if file_list is None:
            file_list = []
        DirectFTPDownload.__init__(self, protocol, host, rootdir, file_list)
        self.save_as = None
        self.method = 'GET'
        self.param = {}

    def download(self, local_dir, keep_dirs=True):
        '''
        Download remote files to local_dir

        :param local_dir: Directory where files should be downloaded
        :type local_dir: str
        :param keep_dirs: keep file name directory structure or copy file in local_dir directly
        :param keep_dirs: bool
        :return: list of downloaded files
        '''
        logging.debug('DirectHTTP:Download')
        nb_files = len(self.files_to_download)

        if nb_files > 1:
            self.files_to_download = []
            logging.error('DirectHTTP accepts only 1 file')

        cur_files = 1

        for rfile in self.files_to_download:
            if self.kill_received:
                raise Exception('Kill request received, exiting')

            if self.save_as is None:
                self.save_as = rfile['name']
            file_dir = local_dir
            if keep_dirs:
                file_dir = local_dir + os.path.dirname(self.save_as)
            file_path = file_dir + '/' + os.path.basename(self.save_as)
            '''
            self.mkdir_lock.acquire()
            try:
                if not os.path.exists(file_dir):
                    os.makedirs(file_dir)
            except Exception as e:
                logging.error(e)
            finally:
                self.mkdir_lock.release() # release lock, no matter what
            '''
            logging.debug('DirectHTTP:Download:Progress'+str(cur_files)+'/'+str(nb_files)+' downloading file '+rfile['name']+', save as '+self.save_as)
            cur_files += 1
            if not 'url' in rfile:
                rfile['url'] = self.url
            fp = open(file_path, "wb")
            curl = pycurl.Curl()

            if self.proxy is not None:
                curl.setopt(pycurl.PROXY, self.proxy)
                if self.proxy_auth is not None:
                    curl.setopt(pycurl.PROXYUSERPWD, self.proxy_auth)

            if self.method == 'POST':
                # Form data must be provided already urlencoded.
                postfields = urllib.parse.urlencode(self.param)
                # Sets request method to POST,
                # Content-Type header to application/x-www-form-urlencoded
                # and data to send in request body.
                if self.credentials is not None:
                    curl.setopt(pycurl.USERPWD, self.credentials)

                curl.setopt(pycurl.POSTFIELDS, postfields)
                try:
                    curl.setopt(pycurl.URL, rfile['url']+rfile['root']+'/'+rfile['name'])
                except Exception as a:
                    curl.setopt(pycurl.URL, (rfile['url']+rfile['root']+'/'+rfile['name']).encode('ascii', 'ignore'))
                #curl.setopt(pycurl.URL, rfile['url']+rfile['root']+'/'+rfile['name'])
            else:
                url = rfile['url']+rfile['root']+'/'+rfile['name']+'?'+urllib.parse.urlencode(self.param)
                #curl.setopt(pycurl.URL, url)
                try:
                    curl.setopt(pycurl.URL, url)
                except Exception as a:
                    curl.setopt(pycurl.URL, url.encode('ascii', 'ignore'))

            curl.setopt(pycurl.WRITEDATA, fp)
            curl.perform()

            curl.close()
            fp.close()
            logging.debug('downloaded!')
            rfile['name'] = self.save_as
            self.set_permissions(file_path, rfile)
            self.set_progress(1, nb_files)
        return self.files_to_download

    def header_function(self, header_line):
        # HTTP standard specifies that headers are encoded in iso-8859-1.
        # On Python 2, decoding step can be skipped.
        # On Python 3, decoding step is required.
        header_line = header_line.decode('iso-8859-1')

        # Header lines include the first status line (HTTP/1.x ...).
        # We are going to ignore all lines that don't have a colon in them.
        # This will botch headers that are split on multiple lines...
        if ':' not in header_line:
            return

        # Break the header line into header name and value.
        name, value = header_line.split(':', 1)

        # Remove whitespace that may be present.
        # Header lines include the trailing newline, and there may be whitespace
        # around the colon.
        name = name.strip()
        value = value.strip()

        # Header names are case insensitive.
        # Lowercase name here.
        name = name.lower()

        # Now we can actually record the header name and value.
        self.headers[name] = value

    def list(self, directory=''):
        '''
        Try to get file headers to get last_modification and size
        '''
        for file in self.files_to_download:
            self.crl.setopt(pycurl.HEADER, True)
            if self.credentials is not None:
                self.crl.setopt(pycurl.USERPWD, self.credentials)

            if self.proxy is not None:
                self.crl.setopt(pycurl.PROXY, self.proxy)
                if self.proxy_auth is not None:
                    self.crl.setopt(pycurl.PROXYUSERPWD, self.proxy_auth)

            self.crl.setopt(pycurl.NOBODY, True)
            try:
                self.crl.setopt(pycurl.URL, self.url+self.rootdir+file['name'])
            except Exception as a:
                self.crl.setopt(pycurl.URL, (self.url+self.rootdir+file['name']).encode('ascii', 'ignore'))
            #self.crl.setopt(pycurl.URL, self.url+self.rootdir+file['name'])
            output = BytesIO()
            # lets assign this buffer to pycurl object
            self.crl.setopt(pycurl.WRITEFUNCTION, output.write)
            self.crl.setopt(pycurl.HEADERFUNCTION, self.header_function)
            self.crl.perform()

            # Figure out what encoding was sent with the response, if any.
            # Check against lowercased header name.
            encoding = None
            if 'content-type' in self.headers:
                content_type = self.headers['content-type'].lower()
                match = re.search('charset=(\S+)', content_type)
                if match:
                    encoding = match.group(1)
            if encoding is None:
                # Default encoding for HTML is iso-8859-1.
                # Other content types may have different default encoding,
                # or in case of binary data, may have no encoding at all.
                encoding = 'iso-8859-1'

            # lets get the output in a string
            result = output.getvalue().decode(encoding)

            lines = re.split(r'[\n\r]+', result)
            for line in lines:
                parts = line.split(':')
                if parts[0].strip() == 'Content-Length':
                    file['size'] = parts[1].strip()
                if parts[0].strip() == 'Last-Modified':
                    # Sun, 06 Nov 1994
                    res = re.match('(\w+),\s+(\d+)\s+(\w+)\s+(\d+)', parts[1].strip())
                    if res:
                        file['hash'] = hashlib.md5(str(res.group(0)).encode('utf-8')).hexdigest()
                        file['day'] = res.group(2)
                        file['month'] = Utils.month_to_num(res.group(3))
                        file['year'] = res.group(4)
                        continue
                    #Sunday, 06-Nov-94
                    res = re.match('(\w+),\s+(\d+)-(\w+)-(\d+)', parts[1].strip())
                    if res:
                        file['hash'] = hashlib.md5(str(res.group(0)).encode('utf-8')).hexdigest()
                        file['day'] = res.group(2)
                        file['month'] = Utils.month_to_num(res.group(3))
                        file['year'] = str(2000 + int(res.group(4)))
                        continue
                    #Sun Nov  6 08:49:37 1994
                    res = re.match('(\w+)\s+(\w+)\s+(\d+)\s+\d{2}:\d{2}:\d{2}\s+(\d+)', parts[1].strip())
                    if res:
                        file['hash'] = hashlib.md5(str(res.group(0)).encode('utf-8')).hexdigest()
                        file['day'] = res.group(3)
                        file['month'] = Utils.month_to_num(res.group(2))
                        file['year'] = res.group(4)
                        continue
        return (self.files_to_download, [])
<commit_msg>add creation dir for tests<commit_after>from future import standard_library
standard_library.install_aliases()
from builtins import str
import datetime
import logging
import pycurl
import io
import os
import re
import urllib.request, urllib.parse, urllib.error
import hashlib

from biomaj.download.interface import DownloadInterface
from biomaj.download.ftp import FTPDownload
from biomaj.utils import Utils

try:
    from io import BytesIO
except ImportError:
    from StringIO import StringIO as BytesIO

class MultiDownload(DownloadInterface):
    '''
    Base interface for a downloader using multiple downloaders
    '''
    def __init__(self):
        DownloadInterface.__init__(self)
        self.downloaders = []
        self.files_to_download = []

    def add_downloaders(self, downloaders):
        '''
        Adds a list of downloaders
        '''
        self.downloaders += downloaders
        for d in downloaders:
            self.files_to_download += d.files_to_download

    def match(self, patterns, file_list, dir_list=None, prefix='', submatch=False):
        if dir_list is None:
            dir_list = []
        self.files_to_download = []
        for d in self.downloaders:
            d.match(patterns, d.files_to_download, [], prefix, submatch)
        self.files_to_download = []
        for d in self.downloaders:
            self.files_to_download += d.files_to_download

    def download(self, local_dir):
        self.files_to_download = []
        for d in self.downloaders:
            if self.kill_received:
                raise Exception('Kill request received, exiting')
            d.download(local_dir)
        self.files_to_download = []
        for d in self.downloaders:
            self.files_to_download += d.files_to_download
        return (self.files_to_download, [])

    def list(self):
        self.files_to_download = []
        for d in self.downloaders:
            d.list()
        self.files_to_download = []
        for d in self.downloaders:
            self.files_to_download += d.files_to_download
        return (self.files_to_download, [])

    def close(self):
        for d in self.downloaders:
            d.close()

class DirectFTPDownload(FTPDownload):
    '''
    download a list of files from FTP, no regexp
    '''

    def __init__(self, protocol, host, rootdir='', file_list=None):
        '''

        Initialize the files in list with today as last-modification date.
        Size is also preset to zero, size will be set after download

        :param file_list: list of files to download on server
        :type file_list: list
        '''
        FTPDownload.__init__(self, protocol, host, rootdir)
        if file_list is None:
            file_list = []
        today = datetime.date.today()
        self.files_to_download = []
        self.headers = {}
        for file in file_list:
            rfile = {}
            rfile['root'] = self.rootdir
            rfile['permissions'] = ''
            rfile['group'] = ''
            rfile['user'] = ''
            rfile['size'] = 0
            rfile['month'] = today.month
            rfile['day'] = today.day
            rfile['year'] = today.year
            rfile['name'] = file
            rfile['hash'] = None
            self.files_to_download.append(rfile)

    def list(self, directory=''):
        '''
        FTP protocol does not give us the possibility to get file date from remote url
        '''
        return (self.files_to_download, [])

    def match(self, patterns, file_list, dir_list=None, prefix='', submatch=False):
        '''
        All files to download match, no pattern
        '''
        if dir_list is None:
            dir_list = []
        self.files_to_download = file_list

class DirectHttpDownload(DirectFTPDownload):

    def __init__(self, protocol, host, rootdir='', file_list=None):
        '''
        :param file_list: list of files to download on server
        :type file_list: list
        '''
        if file_list is None:
            file_list = []
        DirectFTPDownload.__init__(self, protocol, host, rootdir, file_list)
        self.save_as = None
        self.method = 'GET'
        self.param = {}

    def download(self, local_dir, keep_dirs=True):
        '''
        Download remote files to local_dir

        :param local_dir: Directory where files should be downloaded
        :type local_dir: str
        :param keep_dirs: keep file name directory structure or copy file in local_dir directly
        :param keep_dirs: bool
        :return: list of downloaded files
        '''
        logging.debug('DirectHTTP:Download')
        nb_files = len(self.files_to_download)

        if nb_files > 1:
            self.files_to_download = []
            logging.error('DirectHTTP accepts only 1 file')

        cur_files = 1

        for rfile in self.files_to_download:
            if self.kill_received:
                raise Exception('Kill request received, exiting')

            if self.save_as is None:
                self.save_as = rfile['name']
            file_dir = local_dir
            if keep_dirs:
                file_dir = local_dir + os.path.dirname(self.save_as)
            file_path = file_dir + '/' + os.path.basename(self.save_as)

            # For unit tests only, workflow will take in charge directory creation before to avoid thread multi access
            if not os.path.exists(file_dir):
                os.makedirs(file_dir)
            '''
            self.mkdir_lock.acquire()
            try:
                if not os.path.exists(file_dir):
                    os.makedirs(file_dir)
            except Exception as e:
                logging.error(e)
            finally:
                self.mkdir_lock.release() # release lock, no matter what
            '''
            logging.debug('DirectHTTP:Download:Progress'+str(cur_files)+'/'+str(nb_files)+' downloading file '+rfile['name']+', save as '+self.save_as)
            cur_files += 1
            if not 'url' in rfile:
                rfile['url'] = self.url
            fp = open(file_path, "wb")
            curl = pycurl.Curl()

            if self.proxy is not None:
                curl.setopt(pycurl.PROXY, self.proxy)
                if self.proxy_auth is not None:
                    curl.setopt(pycurl.PROXYUSERPWD, self.proxy_auth)

            if self.method == 'POST':
                # Form data must be provided already urlencoded.
                postfields = urllib.parse.urlencode(self.param)
                # Sets request method to POST,
                # Content-Type header to application/x-www-form-urlencoded
                # and data to send in request body.
                if self.credentials is not None:
                    curl.setopt(pycurl.USERPWD, self.credentials)

                curl.setopt(pycurl.POSTFIELDS, postfields)
                try:
                    curl.setopt(pycurl.URL, rfile['url']+rfile['root']+'/'+rfile['name'])
                except Exception as a:
                    curl.setopt(pycurl.URL, (rfile['url']+rfile['root']+'/'+rfile['name']).encode('ascii', 'ignore'))
                #curl.setopt(pycurl.URL, rfile['url']+rfile['root']+'/'+rfile['name'])
            else:
                url = rfile['url']+rfile['root']+'/'+rfile['name']+'?'+urllib.parse.urlencode(self.param)
                #curl.setopt(pycurl.URL, url)
                try:
                    curl.setopt(pycurl.URL, url)
                except Exception as a:
                    curl.setopt(pycurl.URL, url.encode('ascii', 'ignore'))

            curl.setopt(pycurl.WRITEDATA, fp)
            curl.perform()

            curl.close()
            fp.close()
            logging.debug('downloaded!')
            rfile['name'] = self.save_as
            self.set_permissions(file_path, rfile)
            self.set_progress(1, nb_files)
        return self.files_to_download

    def header_function(self, header_line):
        # HTTP standard specifies that headers are encoded in iso-8859-1.
        # On Python 2, decoding step can be skipped.
        # On Python 3, decoding step is required.
        header_line = header_line.decode('iso-8859-1')

        # Header lines include the first status line (HTTP/1.x ...).
        # We are going to ignore all lines that don't have a colon in them.
        # This will botch headers that are split on multiple lines...
        if ':' not in header_line:
            return

        # Break the header line into header name and value.
        name, value = header_line.split(':', 1)

        # Remove whitespace that may be present.
        # Header lines include the trailing newline, and there may be whitespace
        # around the colon.
        name = name.strip()
        value = value.strip()

        # Header names are case insensitive.
        # Lowercase name here.
        name = name.lower()

        # Now we can actually record the header name and value.
        self.headers[name] = value

    def list(self, directory=''):
        '''
        Try to get file headers to get last_modification and size
        '''
        for file in self.files_to_download:
            self.crl.setopt(pycurl.HEADER, True)
            if self.credentials is not None:
                self.crl.setopt(pycurl.USERPWD, self.credentials)

            if self.proxy is not None:
                self.crl.setopt(pycurl.PROXY, self.proxy)
                if self.proxy_auth is not None:
                    self.crl.setopt(pycurl.PROXYUSERPWD, self.proxy_auth)

            self.crl.setopt(pycurl.NOBODY, True)
            try:
                self.crl.setopt(pycurl.URL, self.url+self.rootdir+file['name'])
            except Exception as a:
                self.crl.setopt(pycurl.URL, (self.url+self.rootdir+file['name']).encode('ascii', 'ignore'))
            #self.crl.setopt(pycurl.URL, self.url+self.rootdir+file['name'])
            output = BytesIO()
            # lets assign this buffer to pycurl object
            self.crl.setopt(pycurl.WRITEFUNCTION, output.write)
            self.crl.setopt(pycurl.HEADERFUNCTION, self.header_function)
            self.crl.perform()

            # Figure out what encoding was sent with the response, if any.
            # Check against lowercased header name.
            encoding = None
            if 'content-type' in self.headers:
                content_type = self.headers['content-type'].lower()
                match = re.search('charset=(\S+)', content_type)
                if match:
                    encoding = match.group(1)
            if encoding is None:
                # Default encoding for HTML is iso-8859-1.
                # Other content types may have different default encoding,
                # or in case of binary data, may have no encoding at all.
                encoding = 'iso-8859-1'

            # lets get the output in a string
            result = output.getvalue().decode(encoding)

            lines = re.split(r'[\n\r]+', result)
            for line in lines:
                parts = line.split(':')
                if parts[0].strip() == 'Content-Length':
                    file['size'] = parts[1].strip()
                if parts[0].strip() == 'Last-Modified':
                    # Sun, 06 Nov 1994
                    res = re.match('(\w+),\s+(\d+)\s+(\w+)\s+(\d+)', parts[1].strip())
                    if res:
                        file['hash'] = hashlib.md5(str(res.group(0)).encode('utf-8')).hexdigest()
                        file['day'] = res.group(2)
                        file['month'] = Utils.month_to_num(res.group(3))
                        file['year'] = res.group(4)
                        continue
                    #Sunday, 06-Nov-94
                    res = re.match('(\w+),\s+(\d+)-(\w+)-(\d+)', parts[1].strip())
                    if res:
                        file['hash'] = hashlib.md5(str(res.group(0)).encode('utf-8')).hexdigest()
                        file['day'] = res.group(2)
                        file['month'] = Utils.month_to_num(res.group(3))
                        file['year'] = str(2000 + int(res.group(4)))
                        continue
                    #Sun Nov  6 08:49:37 1994
                    res = re.match('(\w+)\s+(\w+)\s+(\d+)\s+\d{2}:\d{2}:\d{2}\s+(\d+)', parts[1].strip())
                    if res:
                        file['hash'] = hashlib.md5(str(res.group(0)).encode('utf-8')).hexdigest()
                        file['day'] = res.group(3)
                        file['month'] = Utils.month_to_num(res.group(2))
                        file['year'] = res.group(4)
                        continue
        return (self.files_to_download, [])
<|endoftext|>


<commit_before># Uses the XHTML files to extract a table of contents for the US Code.
# run with: ./run xhtml_to_structure year=2011 > structure.json

import glob, re, lxml.html, json, sys
import HTMLParser

pars = HTMLParser.HTMLParser()
section_symbol = u'\xa7'

def run(options):
  # Output data structure.
  TOC = [ ]

  path = None

  # Loop through all titles of the code...
  for fn in glob.glob("data/uscode.house.gov/xhtml/" + options["year"] + "/*.htm"):
    if not re.search(options["year"] + r"usc\d+a?\.htm$", fn): continue # is it a title file?

    # Parse the XHTML file...
    dom = lxml.html.parse(fn)

    # The file structure is flat. Loop through the XML
    # nodes in this title.
    for n in dom.find("body/div"):
      # Look for comments of the form <!-- expcite:... -->
      # This tells us the current table of contents location for the following <h3>.
      m = re.match(ur"<!-- expcite:(.*\S)\s*-->", unicode(n))
      if m:
        # This is a "!@!"-separated string giving the table-of-contents
        # path to each section as we see it.
        expcite = m.group(1)

        # These comments have HTML entities. Replace them with unicode.
        expcite = expcite.replace(" ", " ")
        expcite = pars.unescape(expcite)

        # Parse the table of contents path.
        path = parse_expcite(expcite)

      elif n.tag == "h3":
        # All headings are h3s. Check if it starts with the curly S section symbol.
        h3 = n.text_content()
        if not h3.startswith(section_symbol): continue

        # The most recent expcite path is the TOC location of this section.
        if not path: raise Exception("h3 without path")

        # Insert the section into our TOC structure.
        parse_h3(path, h3, TOC)

        # Clear so we don't reuse the path on the next h3.
        path = None

  # Sort the titles (take into account appendix notation).
  TOC.sort(key = lambda title : (int(title[0][1].replace("a", "")), title[0][1]))

  # Reformat the output.
  TOC = [reformat_structure(title) for title in TOC]

  # Write output in JSON to stdout.
  json.dump(TOC, sys.stdout, indent=True, sort_keys=True, check_circular=False)

def parse_expcite(expcite):
  path = expcite.split("!@!")

  # Parse each part of the path:
  for i in xrange(len(path)):
    if re.match(r"\[.*-(REPEALED|RESERVED|OMITTED|TRANSFERRED)\]\s*$", path[i], re.I):
      # This part is repealed. No need to process this path at all.
      path = None
      break

    m = re.match(r"(TITLE|SUBTITLE|CHAPTER|SUBCHAPTER|PART|SUBPART|DIVISION) ([^\-]+)-(.*)$|Secs?\. (.*)", path[i], re.I)
    if not m:
      # Some random text string. It's a part of a title with no level specifier.
      # We'll call it a "heading" level. In Title 50a, there are just names of
      # acts under the title, no apparent chapter.
      path[i] = ("heading", None, path[i])

    elif m.group(1):
      # Matches TITLE|CHAPTER...
      # Store as (TITLE|CHAPTER, NUMBER, NAME)
      # Replace en-dashes in the number with simple dashes, as we do with section numbers.
      path[i] = (m.group(1).lower(), m.group(2).replace(u"\u2013", "-"), m.group(3))

      # Reformat title appendices: XXX, APPENDIX => XXXa.
      if i == 0 and path[i][0] == "title" and ", APPENDIX" in path[i][1]:
        path[i] = (path[i][0], path[i][1].replace(", APPENDIX", "a"), path[i][2] + " (APPENDIX)")
      elif i == 0 and path[i][0] == "title" and path[i][2] == "APPENDIX": # titles 5, 18 look like this
        path[i] = (path[i][0], path[i][1] + "a", path[i][2])

    elif m.group(4) and i == len(path) - 1:
      # Matches a section number or range of sections.
      # We'll get this information from the next <h3> element, which also has the
      # section title.
      path.pop()

    else:
      raise Exception("Invalid expcite?")

  return path

def parse_h3(path, h3, TOC):
  # Skip sections that are just placeholders.
  if re.match(section_symbol + section_symbol + r"?(.*?)\. (Repealed.*|Transferred|Omitted)(\.|$)", h3):
    # This is for multiple sections, which are always repealed, or
    # repealed/transferred sections.
    return

  # Reformat section numbers. Replace en-dashes with simple dashes, as we do with chapter etc. numbers.
  h3 = h3.replace(u"\u2013", "-")

  # Parse the section number and description, and add that to the path.
  m = re.match(section_symbol + r"(.*?)\.? (.*)", h3)
  if not m: raise Exception("Could not parse: " + h3)
  path.append( ("section", m.group(1), m.group(2)) )

  # Add the new path into the TOC, making a structure like:
  #  [ ( (title, 17, Copyright), [
  #       ... sub parts ..
  #      ])
  #  ]
  _toc = TOC
  for p in path:
    if len(_toc) == 0 or _toc[-1][0] != p:
      _toc.append( (p, []) )
    _toc = _toc[-1][1] # move in

def reformat_structure(entry):
  ret = {
    "level": entry[0][0],
    "number": entry[0][1],
    "name": entry[0][2],
  }
  if len(entry[1]):
    ret["subparts"] = [reformat_structure(e) for e in entry[1]]
  return ret

<commit_msg>default year to uscprelim, update regex check on filename to accommodate<commit_after># Uses the XHTML files to extract a table of contents for the US Code.
# run with: ./run xhtml_to_structure year=2011 > structure.json

import glob, re, lxml.html, json, sys
import HTMLParser

pars = HTMLParser.HTMLParser()
section_symbol = u'\xa7'

def run(options):
  # Output data structure.
  TOC = [ ]

  year = options.get("year", "uscprelim") # default to USCprelim

  path = None

  # Loop through all titles of the code...
  for fn in glob.glob("data/uscode.house.gov/xhtml/" + year + "/*.htm"):
    if year == "uscprelim":
      match = "PRELIM"
    else:
      match = year

    if not re.search(match + r"usc\d+a?\.htm$", fn, re.I): continue # is it a title file?

    # Parse the XHTML file...
    dom = lxml.html.parse(fn)

    # The file structure is flat. Loop through the XML
    # nodes in this title.
    for n in dom.find("body/div"):
      # Look for comments of the form <!-- expcite:... -->
      # This tells us the current table of contents location for the following <h3>.
      m = re.match(ur"<!-- expcite:(.*\S)\s*-->", unicode(n))
      if m:
        # This is a "!@!"-separated string giving the table-of-contents
        # path to each section as we see it.
        expcite = m.group(1)

        # These comments have HTML entities. Replace them with unicode.
        expcite = expcite.replace(" ", " ")
        expcite = pars.unescape(expcite)

        # Parse the table of contents path.
        path = parse_expcite(expcite)

      elif n.tag == "h3":
        # All headings are h3s. Check if it starts with the curly S section symbol.
        h3 = n.text_content()
        if not h3.startswith(section_symbol): continue

        # The most recent expcite path is the TOC location of this section.
        if not path: raise Exception("h3 without path")

        # Insert the section into our TOC structure.
        parse_h3(path, h3, TOC)

        # Clear so we don't reuse the path on the next h3.
        path = None

  # Sort the titles (take into account appendix notation).
  TOC.sort(key = lambda title : (int(title[0][1].replace("a", "")), title[0][1]))

  # Reformat the output.
  TOC = [reformat_structure(title) for title in TOC]

  # Write output in JSON to stdout.
  json.dump(TOC, sys.stdout, indent=True, sort_keys=True, check_circular=False)

def parse_expcite(expcite):
  path = expcite.split("!@!")

  # Parse each part of the path:
  for i in xrange(len(path)):
    if re.match(r"\[.*-(REPEALED|RESERVED|OMITTED|TRANSFERRED)\]\s*$", path[i], re.I):
      # This part is repealed. No need to process this path at all.
      path = None
      break

    m = re.match(r"(TITLE|SUBTITLE|CHAPTER|SUBCHAPTER|PART|SUBPART|DIVISION) ([^\-]+)-(.*)$|Secs?\. (.*)", path[i], re.I)
    if not m:
      # Some random text string. It's a part of a title with no level specifier.
      # We'll call it a "heading" level. In Title 50a, there are just names of
      # acts under the title, no apparent chapter.
      path[i] = ("heading", None, path[i])

    elif m.group(1):
      # Matches TITLE|CHAPTER...
      # Store as (TITLE|CHAPTER, NUMBER, NAME)
      # Replace en-dashes in the number with simple dashes, as we do with section numbers.
      path[i] = (m.group(1).lower(), m.group(2).replace(u"\u2013", "-"), m.group(3))

      # Reformat title appendices: XXX, APPENDIX => XXXa.
      if i == 0 and path[i][0] == "title" and ", APPENDIX" in path[i][1]:
        path[i] = (path[i][0], path[i][1].replace(", APPENDIX", "a"), path[i][2] + " (APPENDIX)")
      elif i == 0 and path[i][0] == "title" and path[i][2] == "APPENDIX": # titles 5, 18 look like this
        path[i] = (path[i][0], path[i][1] + "a", path[i][2])

    elif m.group(4) and i == len(path) - 1:
      # Matches a section number or range of sections.
      # We'll get this information from the next <h3> element, which also has the
      # section title.
      path.pop()

    else:
      raise Exception("Invalid expcite?")

  return path

def parse_h3(path, h3, TOC):
  # Skip sections that are just placeholders.
  if re.match(section_symbol + section_symbol + r"?(.*?)\. (Repealed.*|Transferred|Omitted)(\.|$)", h3):
    # This is for multiple sections, which are always repealed, or
    # repealed/transferred sections.
    return

  # Reformat section numbers. Replace en-dashes with simple dashes, as we do with chapter etc. numbers.
  h3 = h3.replace(u"\u2013", "-")

  # Parse the section number and description, and add that to the path.
  m = re.match(section_symbol + r"(.*?)\.? (.*)", h3)
  if not m: raise Exception("Could not parse: " + h3)
  path.append( ("section", m.group(1), m.group(2)) )

  # Add the new path into the TOC, making a structure like:
  #  [ ( (title, 17, Copyright), [
  #       ... sub parts ..
  #      ])
  #  ]
  _toc = TOC
  for p in path:
    if len(_toc) == 0 or _toc[-1][0] != p:
      _toc.append( (p, []) )
    _toc = _toc[-1][1] # move in

def reformat_structure(entry):
  ret = {
    "level": entry[0][0],
    "number": entry[0][1],
    "name": entry[0][2],
  }
  if len(entry[1]):
    ret["subparts"] = [reformat_structure(e) for e in entry[1]]
  return ret

<|endoftext|>


<commit_before># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import
from __future__ import division
import re

import dataproperty
import ipaddress
import six
from subprocrunner import logger
from subprocrunner import SubprocessRunner

from .parser import TcFilterParser
from .parser import TcQdiscParser
from ._common import verify_network_interface
from ._converter import Humanreadable
from ._error import TcCommandExecutionError

def _validate_within_min_max(param_name, value, min_value, max_value):
    if value is None:
        return

    if value > max_value:
        raise ValueError(
            "{:s} is too high: expected<={:f}[%], value={:f}[%]".format(
                param_name, max_value, value))

    if value < min_value:
        raise ValueError(
            "{:s} is too low: expected>={:f}[%], value={:f}[%]".format(
                param_name, min_value, value))

class TrafficDirection(object):
    OUTGOING = "outgoing"
    INCOMING = "incoming"
    LIST = [OUTGOING, INCOMING]

class TrafficControl(object):
    __OUT_DEVICE_QDISC_MINOR_ID = 1
    __IN_DEVICE_QDISC_MINOR_ID = 3

    __MIN_PACKET_LOSS_RATE = 0  # [%]
    __MAX_PACKET_LOSS_RATE = 99  # [%]

    __MIN_LATENCY_MS = 0  # [millisecond]
    __MAX_LATENCY_MS = 10000  # [millisecond]

    __MIN_CORRUPTION_RATE = 0  # [%]
    __MAX_CORRUPTION_RATE = 99  # [%]

    __MIN_BUFFER_BYTE = 1600

    __MIN_PORT = 0
    __MAX_PORT = 65535

    __REGEXP_FILE_EXISTS = re.compile("RTNETLINK answers: File exists")

    __EXISTS_MSG_TEMPLATE = (
        "{:s} "
        "execute with --overwrite option if you want to overwrite "
        "the existing settings.")

    @property
    def ifb_device(self):
        return "ifb{:d}".format(self.__get_device_qdisc_major_id())

    def __init__(self, device):
        self.__device = device

        self.direction = None
        self.bandwidth_rate = None  # bandwidth string [G/M/K bps]
        self.latency_ms = None  # [milliseconds]
        self.latency_distro_ms = None  # [milliseconds]
        self.packet_loss_rate = None  # [%]
        self.curruption_rate = None  # [%]
        self.network = None
        self.port = None

        self.src_network = None

    def validate(self):
        verify_network_interface(self.__device)
        self.__validate_bandwidth_rate()
        self.__validate_network_delay()
        self.__validate_packet_loss_rate()
        self.__validate_curruption_rate()
        self.network = self.__sanitize_network(self.network)
        self.src_network = self.__sanitize_network(self.src_network)
        self.__validate_port()

    def set_tc(self):
        self.__setup_ifb()

        qdisc_major_id = self.__get_device_qdisc_major_id()
        self.__make_qdisc(qdisc_major_id)
        self.__set_pre_network_filter(qdisc_major_id)
        self.__set_netem(qdisc_major_id)
        self.__set_network_filter(qdisc_major_id)
        self.__set_rate(qdisc_major_id)

    def delete_tc(self):
        returncode = 0

        returncode |= self.__run(
            "tc qdisc del dev {:s} root".format(self.__device),
            re.compile("RTNETLINK answers: No such file or directory"),
            "failed to delete qdisc: no qdisc for outgoing packets")

        returncode |= self.__run(
            "tc qdisc del dev {:s} ingress".format(self.__device),
            re.compile("RTNETLINK answers: Invalid argument"),
            "failed to qdisc: no qdisc for incomming packets")

        returncode |= self.__delete_ifb_device()

        return returncode

    def get_tc_parameter(self):
        return {
            self.__device: {
                TrafficDirection.OUTGOING: self.__get_filter(self.__device),
                TrafficDirection.INCOMING: self.__get_filter(
                    self.__get_ifb_from_device(self.__device)),
            },
        }

    def __validate_bandwidth_rate(self):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return

        rate = Humanreadable().humanreadable_to_byte(self.bandwidth_rate)
        if rate <= 0:
            raise ValueError("rate must be greater than zero")

    def __validate_network_delay(self):
        _validate_within_min_max(
            "latency_ms",
            self.latency_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

        _validate_within_min_max(
            "latency_distro_ms",
            self.latency_distro_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

    def __validate_packet_loss_rate(self):
        _validate_within_min_max(
            "packet_loss_rate",
            self.packet_loss_rate,
            self.__MIN_PACKET_LOSS_RATE, self.__MAX_PACKET_LOSS_RATE)

    def __validate_curruption_rate(self):
        _validate_within_min_max(
            "curruption_rate",
            self.curruption_rate,
            self.__MIN_CORRUPTION_RATE, self.__MAX_CORRUPTION_RATE)

    @staticmethod
    def __sanitize_network(network):
        """
        :return: Network string
        :rtype: str
        :raises ValueError: if the network string is invalid.
        """

        if dataproperty.is_empty_string(network):
            return ""

        try:
            ipaddress.IPv4Address(six.u(network))
            return network + "/32"
        except ipaddress.AddressValueError:
            pass

        ipaddress.IPv4Network(six.u(network))  # validate network str

        return network

    def __validate_port(self):
        _validate_within_min_max(
            "port", self.port, self.__MIN_PORT, self.__MAX_PORT)

    def __make_qdisc(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "root",
            "handle {:x}:".format(qdisc_major_id),
            "prio",
        ]

        return self.__run(
            " ".join(command_list), self.__REGEXP_FILE_EXISTS,
            self.__EXISTS_MSG_TEMPLATE.format(
                "failed to add qdisc: prio qdisc already exists."))

    def __get_tc_device(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__device

        if self.direction == TrafficDirection.INCOMING:
            return self.ifb_device

        raise ValueError("unknown direction: " + self.direction)

    def __get_ifb_from_device(self, device):
        filter_parser = TcFilterParser()
        command = "tc filter show dev {:s} root".format(device)
        filter_runner = SubprocessRunner(command)
        filter_runner.run()

        return filter_parser.parse_incoming_device(filter_runner.stdout)

    def __get_network_direction_str(self):
        if self.direction == TrafficDirection.OUTGOING:
            return "dst"

        if self.direction == TrafficDirection.INCOMING:
            return "src"

        raise ValueError("unknown direction: " + self.direction)

    def __get_device_qdisc_major_id(self):
        import hashlib

        base_device_hash = hashlib.md5(six.b(self.__device)).hexdigest()[:3]
        device_hash_prefix = "1"

        return int(device_hash_prefix + base_device_hash, 16)

    def __get_netem_qdisc_major_id(self, base_qdisc_major_id):
        base_offset = 10

        if self.direction == TrafficDirection.OUTGOING:
            direction_offset = 0
        elif self.direction == TrafficDirection.INCOMING:
            direction_offset = 1

        return base_qdisc_major_id + base_offset + direction_offset

    def __get_qdisc_minor_id(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__OUT_DEVICE_QDISC_MINOR_ID

        if self.direction == TrafficDirection.INCOMING:
            return self.__IN_DEVICE_QDISC_MINOR_ID

        raise ValueError("unknown direction: " + self.direction)

    def __get_filter(self, device):
        if dataproperty.is_empty_string(device):
            return {}

        qdisc_parser = TcQdiscParser()
        filter_parser = TcFilterParser()

        # parse qdisc ---
        command = "tc qdisc show dev {:s}".format(device)
        qdisk_show_runner = SubprocessRunner(command)
        qdisk_show_runner.run()
        qdisc_param = qdisc_parser.parse(qdisk_show_runner.stdout)

        # parse filter ---
        command = "tc filter show dev {:s}".format(device)
        filter_show_runner = SubprocessRunner(command)
        filter_show_runner.run()

        filter_table = {}
        for filter_param in filter_parser.parse_filter(filter_show_runner.stdout):
            key_item_list = []

            if dataproperty.is_not_empty_string(filter_param.get("network")):
                key_item_list.append("network=" + filter_param.get("network"))

            if dataproperty.is_integer(filter_param.get("port")):
                key_item_list.append(
                    "port={:d}".format(filter_param.get("port")))

            filter_key = ", ".join(key_item_list)
            filter_table[filter_key] = {}
            if filter_param.get("flowid") == qdisc_param.get("parent"):
                work_qdisc_param = dict(qdisc_param)
                del work_qdisc_param["parent"]
                filter_table[filter_key] = work_qdisc_param

        return filter_table

    def __setup_ifb(self):
        if self.direction != TrafficDirection.INCOMING:
            return 0

        if dataproperty.is_empty_string(self.ifb_device):
            return -1

        return_code = 0

        command = "modprobe ifb"
        return_code |= SubprocessRunner(command).run()

        command = "ip link add {:s} type ifb".format(self.ifb_device)
        return_code |= SubprocessRunner(command).run()

        command = "ip link set dev {:s} up".format(self.ifb_device)
        return_code |= SubprocessRunner(command).run()

        command = "tc qdisc add dev {:s} ingress".format(self.__device)
        return_code |= SubprocessRunner(command).run()

        command_list = [
            "tc filter add",
            "dev " + self.__device,
            "parent ffff: protocol ip u32 match u32 0 0",
            "flowid {:x}:".format(self.__get_device_qdisc_major_id()),
            "action mirred egress redirect",
            "dev " + self.ifb_device,
        ]
        return_code |= SubprocessRunner(" ".join(command_list)).run()

        return return_code

    def __set_pre_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            not dataproperty.is_integer(self.port),
        ]):
            flowid = "{:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id())
        else:
            flowid = "{:x}:2".format(qdisc_major_id)

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 2 u32 match ip {:s} 0.0.0.0/0".format(
                self.__get_network_direction_str()),
            "flowid " + flowid,
        ]

        if all([
            self.direction == TrafficDirection.OUTGOING,
            dataproperty.is_not_empty_string(self.src_network),
        ]):
            command_list.append("match ip src {:s}".format(self.src_network))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_netem(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
            "handle {:x}:".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id)),
            "netem",
        ]
        if self.packet_loss_rate > 0:
            command_list.append("loss {:f}%".format(self.packet_loss_rate))
        if self.latency_ms > 0:
            command_list.append("delay {:f}ms".format(self.latency_ms))

            if self.latency_distro_ms > 0:
                command_list.append("{:f}ms distribution normal".format(
                    self.latency_distro_ms))

        if self.corruption_rate > 0:
            command_list.append("corrupt {:f}%".format(self.corruption_rate))

        return self.__run(
            " ".join(command_list), self.__REGEXP_FILE_EXISTS,
            self.__EXISTS_MSG_TEMPLATE.format(
                "failed to add qdisc: netem qdisc already exists."))

    def __set_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            self.port is None,
        ]):
            return 0

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 1 u32",
            "flowid {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
        ]
        if dataproperty.is_not_empty_string(self.network):
            command_list.append("match ip {:s} {:s}".format(
                self.__get_network_direction_str(), self.network))
        if self.port is not None:
            command_list.append("match ip dport {:d} 0xffff".format(self.port))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_rate(self, qdisc_major_id):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return 0

        rate_kbps = Humanreadable(kilo_size=1000).humanreadable_to_byte(
            self.bandwidth_rate) / 1000.0
        if rate_kbps <= 0:
            raise ValueError("rate must be greater than zero")

        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id),
                self.__get_qdisc_minor_id()),
            "handle 20:",
            "tbf",
            "rate {:f}kbit".format(rate_kbps),
            "buffer {:d}".format(
                max(rate_kbps, self.__MIN_BUFFER_BYTE)),  # [byte]
            "limit 10000",
        ]

        return SubprocessRunner(" ".join(command_list)).run()

    def __delete_ifb_device(self):
        verify_network_interface(self.ifb_device)

        command_list = [
            "tc qdisc del dev {:s} root".format(self.ifb_device),
            "ip link set dev {:s} down".format(self.ifb_device),
            "ip link delete {:s} type ifb".format(self.ifb_device),
        ]

        if all([
            SubprocessRunner(command).run() != 0
            for command in command_list
        ]):
            return 2

        return 0

    @staticmethod
    def __run(command, regexp, message):
        proc = SubprocessRunner(command, regexp)
        if proc.run() == 0:
            return 0

        match = regexp.search(proc.stderr)
        if match is None:
            return proc.returncode

        logger.notice(message)

        return proc.returncode
<commit_msg>Modify log messages<commit_after># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import
from __future__ import division
import re

import dataproperty
import ipaddress
import six
from subprocrunner import logger
from subprocrunner import SubprocessRunner

from .parser import TcFilterParser
from .parser import TcQdiscParser
from ._common import verify_network_interface
from ._converter import Humanreadable
from ._error import TcCommandExecutionError

def _validate_within_min_max(param_name, value, min_value, max_value):
    if value is None:
        return

    if value > max_value:
        raise ValueError(
            "{:s} is too high: expected<={:f}[%], value={:f}[%]".format(
                param_name, max_value, value))

    if value < min_value:
        raise ValueError(
            "{:s} is too low: expected>={:f}[%], value={:f}[%]".format(
                param_name, min_value, value))

class TrafficDirection(object):
    OUTGOING = "outgoing"
    INCOMING = "incoming"
    LIST = [OUTGOING, INCOMING]

class TrafficControl(object):
    __OUT_DEVICE_QDISC_MINOR_ID = 1
    __IN_DEVICE_QDISC_MINOR_ID = 3

    __MIN_PACKET_LOSS_RATE = 0  # [%]
    __MAX_PACKET_LOSS_RATE = 99  # [%]

    __MIN_LATENCY_MS = 0  # [millisecond]
    __MAX_LATENCY_MS = 10000  # [millisecond]

    __MIN_CORRUPTION_RATE = 0  # [%]
    __MAX_CORRUPTION_RATE = 99  # [%]

    __MIN_BUFFER_BYTE = 1600

    __MIN_PORT = 0
    __MAX_PORT = 65535

    __REGEXP_FILE_EXISTS = re.compile("RTNETLINK answers: File exists")

    __EXISTS_MSG_TEMPLATE = (
        "{:s} "
        "execute with --overwrite option if you want to overwrite "
        "the existing settings.")

    @property
    def ifb_device(self):
        return "ifb{:d}".format(self.__get_device_qdisc_major_id())

    def __init__(self, device):
        self.__device = device

        self.direction = None
        self.bandwidth_rate = None  # bandwidth string [G/M/K bps]
        self.latency_ms = None  # [milliseconds]
        self.latency_distro_ms = None  # [milliseconds]
        self.packet_loss_rate = None  # [%]
        self.curruption_rate = None  # [%]
        self.network = None
        self.port = None

        self.src_network = None

    def validate(self):
        verify_network_interface(self.__device)
        self.__validate_bandwidth_rate()
        self.__validate_network_delay()
        self.__validate_packet_loss_rate()
        self.__validate_curruption_rate()
        self.network = self.__sanitize_network(self.network)
        self.src_network = self.__sanitize_network(self.src_network)
        self.__validate_port()

    def set_tc(self):
        self.__setup_ifb()

        qdisc_major_id = self.__get_device_qdisc_major_id()
        self.__make_qdisc(qdisc_major_id)
        self.__set_pre_network_filter(qdisc_major_id)
        self.__set_netem(qdisc_major_id)
        self.__set_network_filter(qdisc_major_id)
        self.__set_rate(qdisc_major_id)

    def delete_tc(self):
        returncode = 0

        returncode |= self.__run(
            "tc qdisc del dev {:s} root".format(self.__device),
            re.compile("RTNETLINK answers: No such file or directory"),
            "failed to delete qdisc: no qdisc for outgoing packets")

        returncode |= self.__run(
            "tc qdisc del dev {:s} ingress".format(self.__device),
            re.compile("RTNETLINK answers: Invalid argument"),
            "failed to qdisc: no qdisc for incomming packets")

        returncode |= self.__delete_ifb_device()

        return returncode

    def get_tc_parameter(self):
        return {
            self.__device: {
                TrafficDirection.OUTGOING: self.__get_filter(self.__device),
                TrafficDirection.INCOMING: self.__get_filter(
                    self.__get_ifb_from_device(self.__device)),
            },
        }

    def __validate_bandwidth_rate(self):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return

        rate = Humanreadable().humanreadable_to_byte(self.bandwidth_rate)
        if rate <= 0:
            raise ValueError("rate must be greater than zero")

    def __validate_network_delay(self):
        _validate_within_min_max(
            "latency_ms",
            self.latency_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

        _validate_within_min_max(
            "latency_distro_ms",
            self.latency_distro_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

    def __validate_packet_loss_rate(self):
        _validate_within_min_max(
            "packet_loss_rate",
            self.packet_loss_rate,
            self.__MIN_PACKET_LOSS_RATE, self.__MAX_PACKET_LOSS_RATE)

    def __validate_curruption_rate(self):
        _validate_within_min_max(
            "curruption_rate",
            self.curruption_rate,
            self.__MIN_CORRUPTION_RATE, self.__MAX_CORRUPTION_RATE)

    @staticmethod
    def __sanitize_network(network):
        """
        :return: Network string
        :rtype: str
        :raises ValueError: if the network string is invalid.
        """

        if dataproperty.is_empty_string(network):
            return ""

        try:
            ipaddress.IPv4Address(six.u(network))
            return network + "/32"
        except ipaddress.AddressValueError:
            pass

        ipaddress.IPv4Network(six.u(network))  # validate network str

        return network

    def __validate_port(self):
        _validate_within_min_max(
            "port", self.port, self.__MIN_PORT, self.__MAX_PORT)

    def __make_qdisc(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "root",
            "handle {:x}:".format(qdisc_major_id),
            "prio",
        ]

        return self.__run(
            " ".join(command_list), self.__REGEXP_FILE_EXISTS,
            self.__EXISTS_MSG_TEMPLATE.format(
                "failed to add qdisc: prio qdisc already exists."))

    def __get_tc_device(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__device

        if self.direction == TrafficDirection.INCOMING:
            return self.ifb_device

        raise ValueError("unknown direction: " + self.direction)

    def __get_ifb_from_device(self, device):
        filter_parser = TcFilterParser()
        command = "tc filter show dev {:s} root".format(device)
        filter_runner = SubprocessRunner(command)
        filter_runner.run()

        return filter_parser.parse_incoming_device(filter_runner.stdout)

    def __get_network_direction_str(self):
        if self.direction == TrafficDirection.OUTGOING:
            return "dst"

        if self.direction == TrafficDirection.INCOMING:
            return "src"

        raise ValueError("unknown direction: " + self.direction)

    def __get_device_qdisc_major_id(self):
        import hashlib

        base_device_hash = hashlib.md5(six.b(self.__device)).hexdigest()[:3]
        device_hash_prefix = "1"

        return int(device_hash_prefix + base_device_hash, 16)

    def __get_netem_qdisc_major_id(self, base_qdisc_major_id):
        base_offset = 10

        if self.direction == TrafficDirection.OUTGOING:
            direction_offset = 0
        elif self.direction == TrafficDirection.INCOMING:
            direction_offset = 1

        return base_qdisc_major_id + base_offset + direction_offset

    def __get_qdisc_minor_id(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__OUT_DEVICE_QDISC_MINOR_ID

        if self.direction == TrafficDirection.INCOMING:
            return self.__IN_DEVICE_QDISC_MINOR_ID

        raise ValueError("unknown direction: " + self.direction)

    def __get_filter(self, device):
        if dataproperty.is_empty_string(device):
            return {}

        qdisc_parser = TcQdiscParser()
        filter_parser = TcFilterParser()

        # parse qdisc ---
        command = "tc qdisc show dev {:s}".format(device)
        qdisk_show_runner = SubprocessRunner(command)
        qdisk_show_runner.run()
        qdisc_param = qdisc_parser.parse(qdisk_show_runner.stdout)

        # parse filter ---
        command = "tc filter show dev {:s}".format(device)
        filter_show_runner = SubprocessRunner(command)
        filter_show_runner.run()

        filter_table = {}
        for filter_param in filter_parser.parse_filter(filter_show_runner.stdout):
            key_item_list = []

            if dataproperty.is_not_empty_string(filter_param.get("network")):
                key_item_list.append("network=" + filter_param.get("network"))

            if dataproperty.is_integer(filter_param.get("port")):
                key_item_list.append(
                    "port={:d}".format(filter_param.get("port")))

            filter_key = ", ".join(key_item_list)
            filter_table[filter_key] = {}
            if filter_param.get("flowid") == qdisc_param.get("parent"):
                work_qdisc_param = dict(qdisc_param)
                del work_qdisc_param["parent"]
                filter_table[filter_key] = work_qdisc_param

        return filter_table

    def __setup_ifb(self):
        if self.direction != TrafficDirection.INCOMING:
            return 0

        if dataproperty.is_empty_string(self.ifb_device):
            return -1

        return_code = 0

        command = "modprobe ifb"
        return_code |= SubprocessRunner(command).run()

        return_code |= self.__run(
            "ip link add {:s} type ifb".format(self.ifb_device),
            self.__REGEXP_FILE_EXISTS,
            self.__EXISTS_MSG_TEMPLATE.format(
                "failed to add ip link: ip link already exists."))

        command = "ip link set dev {:s} up".format(self.ifb_device)
        return_code |= SubprocessRunner(command).run()

        return_code |= self.__run(
            "tc qdisc add dev {:s} ingress".format(self.__device),
            self.__REGEXP_FILE_EXISTS,
            self.__EXISTS_MSG_TEMPLATE.format(
                "failed to add qdisc: ingress qdisc already exists."))

        command_list = [
            "tc filter add",
            "dev " + self.__device,
            "parent ffff: protocol ip u32 match u32 0 0",
            "flowid {:x}:".format(self.__get_device_qdisc_major_id()),
            "action mirred egress redirect",
            "dev " + self.ifb_device,
        ]
        return_code |= SubprocessRunner(" ".join(command_list)).run()

        return return_code

    def __set_pre_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            not dataproperty.is_integer(self.port),
        ]):
            flowid = "{:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id())
        else:
            flowid = "{:x}:2".format(qdisc_major_id)

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 2 u32 match ip {:s} 0.0.0.0/0".format(
                self.__get_network_direction_str()),
            "flowid " + flowid,
        ]

        if all([
            self.direction == TrafficDirection.OUTGOING,
            dataproperty.is_not_empty_string(self.src_network),
        ]):
            command_list.append("match ip src {:s}".format(self.src_network))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_netem(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
            "handle {:x}:".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id)),
            "netem",
        ]
        if self.packet_loss_rate > 0:
            command_list.append("loss {:f}%".format(self.packet_loss_rate))
        if self.latency_ms > 0:
            command_list.append("delay {:f}ms".format(self.latency_ms))

            if self.latency_distro_ms > 0:
                command_list.append("{:f}ms distribution normal".format(
                    self.latency_distro_ms))

        if self.corruption_rate > 0:
            command_list.append("corrupt {:f}%".format(self.corruption_rate))

        return self.__run(
            " ".join(command_list), self.__REGEXP_FILE_EXISTS,
            self.__EXISTS_MSG_TEMPLATE.format(
                "failed to add qdisc: netem qdisc already exists."))

    def __set_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            self.port is None,
        ]):
            return 0

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 1 u32",
            "flowid {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
        ]
        if dataproperty.is_not_empty_string(self.network):
            command_list.append("match ip {:s} {:s}".format(
                self.__get_network_direction_str(), self.network))
        if self.port is not None:
            command_list.append("match ip dport {:d} 0xffff".format(self.port))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_rate(self, qdisc_major_id):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return 0

        rate_kbps = Humanreadable(kilo_size=1000).humanreadable_to_byte(
            self.bandwidth_rate) / 1000.0
        if rate_kbps <= 0:
            raise ValueError("rate must be greater than zero")

        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id),
                self.__get_qdisc_minor_id()),
            "handle 20:",
            "tbf",
            "rate {:f}kbit".format(rate_kbps),
            "buffer {:d}".format(
                max(rate_kbps, self.__MIN_BUFFER_BYTE)),  # [byte]
            "limit 10000",
        ]

        return SubprocessRunner(" ".join(command_list)).run()

    def __delete_ifb_device(self):
        verify_network_interface(self.ifb_device)

        command_list = [
            "tc qdisc del dev {:s} root".format(self.ifb_device),
            "ip link set dev {:s} down".format(self.ifb_device),
            "ip link delete {:s} type ifb".format(self.ifb_device),
        ]

        if all([
            SubprocessRunner(command).run() != 0
            for command in command_list
        ]):
            return 2

        return 0

    @staticmethod
    def __run(command, regexp, message):
        proc = SubprocessRunner(command, regexp)
        if proc.run() == 0:
            return 0

        match = regexp.search(proc.stderr)
        if match is None:
            return proc.returncode

        logger.notice(message)

        return proc.returncode
<|endoftext|>


<commit_before># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import
from __future__ import division

import dataproperty
import ipaddress
import six
from subprocrunner import SubprocessRunner

from .parser import TcFilterParser
from .parser import TcQdiscParser
from ._common import verify_network_interface
from ._converter import Humanreadable
from ._error import TcCommandExecutionError

def _validate_within_min_max(param_name, value, min_value, max_value):
    if value is None:
        return

    if value > max_value:
        raise ValueError(
            "{:s} is too high: expected<={:f}[%], value={:f}[%]".format(
                param_name, max_value, value))

    if value < min_value:
        raise ValueError(
            "{:s} is too low: expected>={:f}[%], value={:f}[%]".format(
                param_name, min_value, value))

class TrafficDirection(object):
    OUTGOING = "outgoing"
    INCOMING = "incoming"
    LIST = [OUTGOING, INCOMING]

class TrafficControl(object):
    __OUT_DEVICE_QDISC_MINOR_ID = 1
    __IN_DEVICE_QDISC_MINOR_ID = 3

    __MIN_PACKET_LOSS_RATE = 0  # [%]
    __MAX_PACKET_LOSS_RATE = 99  # [%]

    __MIN_LATENCY_MS = 0  # [millisecond]
    __MAX_LATENCY_MS = 10000  # [millisecond]

    __MIN_CORRUPTION_RATE = 0  # [%]
    __MAX_CORRUPTION_RATE = 99  # [%]

    __MIN_BUFFER_BYTE = 1600

    __MIN_PORT = 0
    __MAX_PORT = 65535

    @property
    def ifb_device(self):
        return "ifb{:d}".format(self.__get_device_qdisc_major_id())

    def __init__(self, device):
        self.__device = device

        self.direction = None
        self.bandwidth_rate = None  # bandwidth string [G/M/K bps]
        self.latency_ms = None  # [milliseconds]
        self.latency_distro_ms = None  # [milliseconds]
        self.packet_loss_rate = None  # [%]
        self.curruption_rate = None  # [%]
        self.network = None
        self.port = None

        self.src_network = None

    def validate(self):
        verify_network_interface(self.__device)
        self.__validate_bandwidth_rate()
        self.__validate_network_delay()
        self.__validate_packet_loss_rate()
        self.__validate_curruption_rate()
        self.network = self.__sanitize_network(self.network)
        self.src_network = self.__sanitize_network(self.src_network)
        self.__validate_port()

    def __get_device_qdisc_major_id(self):
        import hashlib

        base_device_hash = hashlib.md5(six.b(self.__device)).hexdigest()[:3]
        device_hash_prefix = "1"

        return int(device_hash_prefix + base_device_hash, 16)

    def set_tc(self):
        self.__setup_ifb()

        qdisc_major_id = self.__get_device_qdisc_major_id()
        self.__make_qdisc(qdisc_major_id)
        self.__set_pre_network_filter(qdisc_major_id)
        self.__set_netem(qdisc_major_id)
        self.__set_network_filter(qdisc_major_id)
        self.__set_rate(qdisc_major_id)

    def delete_tc(self):
        return_code_list = []
        command_list = [
            "tc qdisc del dev {:s} root".format(self.__device),
            "tc qdisc del dev {:s} ingress".format(self.__device),
            "tc qdisc del dev {:s} root".format(self.ifb_device),
            "ip link set dev {:s} down".format(self.ifb_device),
            "ip link delete {:s} type ifb".format(self.ifb_device),
        ]

        for command in command_list:
            return_code_list.append(SubprocessRunner(command).run() != 0)

        return -1 if all(return_code_list) else 0

    def get_tc_parameter(self):
        return {
            self.__device: {
                TrafficDirection.OUTGOING: self.__get_filter(self.__device),
                TrafficDirection.INCOMING: self.__get_filter(
                    self.__get_ifb_from_device(self.__device)),
            },
        }

    def __setup_ifb(self):
        if self.direction != TrafficDirection.INCOMING:
            return

        if dataproperty.is_empty_string(self.ifb_device):
            return

        return_code = 0

        command = "modprobe ifb"
        return_code |= SubprocessRunner(command).run()

        command = "ip link add {:s} type ifb".format(self.ifb_device)
        return_code |= SubprocessRunner(command).run()

        command = "ip link set dev {:s} up".format(self.ifb_device)
        return_code |= SubprocessRunner(command).run()

        command = "tc qdisc add dev {:s} ingress".format(self.__device)
        return_code |= SubprocessRunner(command).run()

        command_list = [
            "tc filter add",
            "dev " + self.__device,
            "parent ffff: protocol ip u32 match u32 0 0",
            "flowid {:x}:".format(self.__get_device_qdisc_major_id()),
            "action mirred egress redirect",
            "dev " + self.ifb_device,
        ]
        return_code |= SubprocessRunner(" ".join(command_list)).run()

        return return_code

    def __validate_bandwidth_rate(self):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return

        rate = Humanreadable().humanreadable_to_byte(self.bandwidth_rate)
        if rate <= 0:
            raise ValueError("rate must be greater than zero")

    def __validate_network_delay(self):
        _validate_within_min_max(
            "latency_ms",
            self.latency_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

        _validate_within_min_max(
            "latency_distro_ms",
            self.latency_distro_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

    def __validate_packet_loss_rate(self):
        _validate_within_min_max(
            "packet_loss_rate",
            self.packet_loss_rate,
            self.__MIN_PACKET_LOSS_RATE, self.__MAX_PACKET_LOSS_RATE)

    def __validate_curruption_rate(self):
        _validate_within_min_max(
            "curruption_rate",
            self.curruption_rate,
            self.__MIN_CORRUPTION_RATE, self.__MAX_CORRUPTION_RATE)

    @staticmethod
    def __sanitize_network(network):
        """
        :return: Network string
        :rtype: str
        :raises ValueError: if the network string is invalid.
        """

        if dataproperty.is_empty_string(network):
            return ""

        try:
            ipaddress.IPv4Address(six.u(network))
            return network + "/32"
        except ipaddress.AddressValueError:
            pass

        ipaddress.IPv4Network(six.u(network))  # validate network str

        return network

    def __validate_port(self):
        _validate_within_min_max(
            "port", self.port, self.__MIN_PORT, self.__MAX_PORT)

    def __make_qdisc(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "root",
            "handle {:x}:".format(qdisc_major_id),
            "prio",
        ]

        return SubprocessRunner(" ".join(command_list)).run()

    def __get_tc_device(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__device

        if self.direction == TrafficDirection.INCOMING:
            return self.ifb_device

        raise ValueError("unknown direction: " + self.direction)

    def __get_ifb_from_device(self, device):
        filter_parser = TcFilterParser()
        command = "tc filter show dev {:s} root".format(device)
        filter_runner = SubprocessRunner(command)
        filter_runner.run()

        return filter_parser.parse_incoming_device(filter_runner.stdout)

    def __get_network_direction_str(self):
        if self.direction == TrafficDirection.OUTGOING:
            return "dst"

        if self.direction == TrafficDirection.INCOMING:
            return "src"

        raise ValueError("unknown direction: " + self.direction)

    def __get_netem_qdisc_major_id(self, base_qdisc_major_id):
        base_offset = 10

        if self.direction == TrafficDirection.OUTGOING:
            direction_offset = 0
        elif self.direction == TrafficDirection.INCOMING:
            direction_offset = 1

        return base_qdisc_major_id + base_offset + direction_offset

    def __get_qdisc_minor_id(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__OUT_DEVICE_QDISC_MINOR_ID

        if self.direction == TrafficDirection.INCOMING:
            return self.__IN_DEVICE_QDISC_MINOR_ID

        raise ValueError("unknown direction: " + self.direction)

    def __get_filter(self, device):
        if dataproperty.is_empty_string(device):
            return {}

        qdisc_parser = TcQdiscParser()
        filter_parser = TcFilterParser()

        # parse qdisc ---
        command = "tc qdisc show dev {:s}".format(device)
        qdisk_show_runner = SubprocessRunner(command)
        qdisk_show_runner.run()
        qdisc_param = qdisc_parser.parse(qdisk_show_runner.stdout)

        # parse filter ---
        command = "tc filter show dev {:s}".format(device)
        filter_show_runner = SubprocessRunner(command)
        filter_show_runner.run()

        filter_table = {}
        for filter_param in filter_parser.parse_filter(filter_show_runner.stdout):
            key_item_list = []

            if dataproperty.is_not_empty_string(filter_param.get("network")):
                key_item_list.append("network=" + filter_param.get("network"))

            if dataproperty.is_integer(filter_param.get("port")):
                key_item_list.append(
                    "port={:d}".format(filter_param.get("port")))

            filter_key = ", ".join(key_item_list)
            filter_table[filter_key] = {}
            if filter_param.get("flowid") == qdisc_param.get("parent"):
                work_qdisc_param = dict(qdisc_param)
                del work_qdisc_param["parent"]
                filter_table[filter_key] = work_qdisc_param

        return filter_table

    def __set_pre_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            not dataproperty.is_integer(self.port),
        ]):
            flowid = "{:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id())
        else:
            flowid = "{:x}:2".format(qdisc_major_id)

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 2 u32 match ip {:s} 0.0.0.0/0".format(
                self.__get_network_direction_str()),
            "flowid " + flowid,
        ]

        if all([
            self.direction == TrafficDirection.OUTGOING,
            dataproperty.is_not_empty_string(self.src_network),
        ]):
            command_list.append("match ip src {:s}".format(self.src_network))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_netem(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
            "handle {:x}:".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id)),
            "netem",
        ]
        if self.packet_loss_rate > 0:
            command_list.append("loss {:f}%".format(self.packet_loss_rate))
        if self.latency_ms > 0:
            command_list.append("delay {:f}ms".format(self.latency_ms))

            if self.latency_distro_ms > 0:
                command_list.append("{:f}ms distribution normal".format(
                    self.latency_distro_ms))

        if self.corruption_rate > 0:
            command_list.append("corrupt {:f}%".format(self.corruption_rate))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            self.port is None,
        ]):
            return 0

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 1 u32",
            "flowid {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
        ]
        if dataproperty.is_not_empty_string(self.network):
            command_list.append("match ip {:s} {:s}".format(
                self.__get_network_direction_str(), self.network))
        if self.port is not None:
            command_list.append("match ip dport {:d} 0xffff".format(self.port))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_rate(self, qdisc_major_id):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return 0

        rate_kbps = Humanreadable(kilo_size=1000).humanreadable_to_byte(
            self.bandwidth_rate) / 1000.0
        if rate_kbps <= 0:
            raise ValueError("rate must be greater than zero")

        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id),
                self.__get_qdisc_minor_id()),
            "handle 20:",
            "tbf",
            "rate {:f}kbit".format(rate_kbps),
            "buffer {:d}".format(
                max(rate_kbps, self.__MIN_BUFFER_BYTE)),  # [byte]
            "limit 10000",
        ]

        return SubprocessRunner(" ".join(command_list)).run()
<commit_msg>Modify for loop to comprehension<commit_after># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import
from __future__ import division

import dataproperty
import ipaddress
import six
from subprocrunner import SubprocessRunner

from .parser import TcFilterParser
from .parser import TcQdiscParser
from ._common import verify_network_interface
from ._converter import Humanreadable
from ._error import TcCommandExecutionError

def _validate_within_min_max(param_name, value, min_value, max_value):
    if value is None:
        return

    if value > max_value:
        raise ValueError(
            "{:s} is too high: expected<={:f}[%], value={:f}[%]".format(
                param_name, max_value, value))

    if value < min_value:
        raise ValueError(
            "{:s} is too low: expected>={:f}[%], value={:f}[%]".format(
                param_name, min_value, value))

class TrafficDirection(object):
    OUTGOING = "outgoing"
    INCOMING = "incoming"
    LIST = [OUTGOING, INCOMING]

class TrafficControl(object):
    __OUT_DEVICE_QDISC_MINOR_ID = 1
    __IN_DEVICE_QDISC_MINOR_ID = 3

    __MIN_PACKET_LOSS_RATE = 0  # [%]
    __MAX_PACKET_LOSS_RATE = 99  # [%]

    __MIN_LATENCY_MS = 0  # [millisecond]
    __MAX_LATENCY_MS = 10000  # [millisecond]

    __MIN_CORRUPTION_RATE = 0  # [%]
    __MAX_CORRUPTION_RATE = 99  # [%]

    __MIN_BUFFER_BYTE = 1600

    __MIN_PORT = 0
    __MAX_PORT = 65535

    @property
    def ifb_device(self):
        return "ifb{:d}".format(self.__get_device_qdisc_major_id())

    def __init__(self, device):
        self.__device = device

        self.direction = None
        self.bandwidth_rate = None  # bandwidth string [G/M/K bps]
        self.latency_ms = None  # [milliseconds]
        self.latency_distro_ms = None  # [milliseconds]
        self.packet_loss_rate = None  # [%]
        self.curruption_rate = None  # [%]
        self.network = None
        self.port = None

        self.src_network = None

    def validate(self):
        verify_network_interface(self.__device)
        self.__validate_bandwidth_rate()
        self.__validate_network_delay()
        self.__validate_packet_loss_rate()
        self.__validate_curruption_rate()
        self.network = self.__sanitize_network(self.network)
        self.src_network = self.__sanitize_network(self.src_network)
        self.__validate_port()

    def __get_device_qdisc_major_id(self):
        import hashlib

        base_device_hash = hashlib.md5(six.b(self.__device)).hexdigest()[:3]
        device_hash_prefix = "1"

        return int(device_hash_prefix + base_device_hash, 16)

    def set_tc(self):
        self.__setup_ifb()

        qdisc_major_id = self.__get_device_qdisc_major_id()
        self.__make_qdisc(qdisc_major_id)
        self.__set_pre_network_filter(qdisc_major_id)
        self.__set_netem(qdisc_major_id)
        self.__set_network_filter(qdisc_major_id)
        self.__set_rate(qdisc_major_id)

    def delete_tc(self):
        command_list = [
            "tc qdisc del dev {:s} root".format(self.__device),
            "tc qdisc del dev {:s} ingress".format(self.__device),
            "tc qdisc del dev {:s} root".format(self.ifb_device),
            "ip link set dev {:s} down".format(self.ifb_device),
            "ip link delete {:s} type ifb".format(self.ifb_device),
        ]

        return_code_list = [
            SubprocessRunner(command).run() != 0
            for command in command_list
        ]

        return -1 if all(return_code_list) else 0

    def get_tc_parameter(self):
        return {
            self.__device: {
                TrafficDirection.OUTGOING: self.__get_filter(self.__device),
                TrafficDirection.INCOMING: self.__get_filter(
                    self.__get_ifb_from_device(self.__device)),
            },
        }

    def __setup_ifb(self):
        if self.direction != TrafficDirection.INCOMING:
            return

        if dataproperty.is_empty_string(self.ifb_device):
            return

        return_code = 0

        command = "modprobe ifb"
        return_code |= SubprocessRunner(command).run()

        command = "ip link add {:s} type ifb".format(self.ifb_device)
        return_code |= SubprocessRunner(command).run()

        command = "ip link set dev {:s} up".format(self.ifb_device)
        return_code |= SubprocessRunner(command).run()

        command = "tc qdisc add dev {:s} ingress".format(self.__device)
        return_code |= SubprocessRunner(command).run()

        command_list = [
            "tc filter add",
            "dev " + self.__device,
            "parent ffff: protocol ip u32 match u32 0 0",
            "flowid {:x}:".format(self.__get_device_qdisc_major_id()),
            "action mirred egress redirect",
            "dev " + self.ifb_device,
        ]
        return_code |= SubprocessRunner(" ".join(command_list)).run()

        return return_code

    def __validate_bandwidth_rate(self):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return

        rate = Humanreadable().humanreadable_to_byte(self.bandwidth_rate)
        if rate <= 0:
            raise ValueError("rate must be greater than zero")

    def __validate_network_delay(self):
        _validate_within_min_max(
            "latency_ms",
            self.latency_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

        _validate_within_min_max(
            "latency_distro_ms",
            self.latency_distro_ms,
            self.__MIN_LATENCY_MS, self.__MAX_LATENCY_MS)

    def __validate_packet_loss_rate(self):
        _validate_within_min_max(
            "packet_loss_rate",
            self.packet_loss_rate,
            self.__MIN_PACKET_LOSS_RATE, self.__MAX_PACKET_LOSS_RATE)

    def __validate_curruption_rate(self):
        _validate_within_min_max(
            "curruption_rate",
            self.curruption_rate,
            self.__MIN_CORRUPTION_RATE, self.__MAX_CORRUPTION_RATE)

    @staticmethod
    def __sanitize_network(network):
        """
        :return: Network string
        :rtype: str
        :raises ValueError: if the network string is invalid.
        """

        if dataproperty.is_empty_string(network):
            return ""

        try:
            ipaddress.IPv4Address(six.u(network))
            return network + "/32"
        except ipaddress.AddressValueError:
            pass

        ipaddress.IPv4Network(six.u(network))  # validate network str

        return network

    def __validate_port(self):
        _validate_within_min_max(
            "port", self.port, self.__MIN_PORT, self.__MAX_PORT)

    def __make_qdisc(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "root",
            "handle {:x}:".format(qdisc_major_id),
            "prio",
        ]

        return SubprocessRunner(" ".join(command_list)).run()

    def __get_tc_device(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__device

        if self.direction == TrafficDirection.INCOMING:
            return self.ifb_device

        raise ValueError("unknown direction: " + self.direction)

    def __get_ifb_from_device(self, device):
        filter_parser = TcFilterParser()
        command = "tc filter show dev {:s} root".format(device)
        filter_runner = SubprocessRunner(command)
        filter_runner.run()

        return filter_parser.parse_incoming_device(filter_runner.stdout)

    def __get_network_direction_str(self):
        if self.direction == TrafficDirection.OUTGOING:
            return "dst"

        if self.direction == TrafficDirection.INCOMING:
            return "src"

        raise ValueError("unknown direction: " + self.direction)

    def __get_netem_qdisc_major_id(self, base_qdisc_major_id):
        base_offset = 10

        if self.direction == TrafficDirection.OUTGOING:
            direction_offset = 0
        elif self.direction == TrafficDirection.INCOMING:
            direction_offset = 1

        return base_qdisc_major_id + base_offset + direction_offset

    def __get_qdisc_minor_id(self):
        if self.direction == TrafficDirection.OUTGOING:
            return self.__OUT_DEVICE_QDISC_MINOR_ID

        if self.direction == TrafficDirection.INCOMING:
            return self.__IN_DEVICE_QDISC_MINOR_ID

        raise ValueError("unknown direction: " + self.direction)

    def __get_filter(self, device):
        if dataproperty.is_empty_string(device):
            return {}

        qdisc_parser = TcQdiscParser()
        filter_parser = TcFilterParser()

        # parse qdisc ---
        command = "tc qdisc show dev {:s}".format(device)
        qdisk_show_runner = SubprocessRunner(command)
        qdisk_show_runner.run()
        qdisc_param = qdisc_parser.parse(qdisk_show_runner.stdout)

        # parse filter ---
        command = "tc filter show dev {:s}".format(device)
        filter_show_runner = SubprocessRunner(command)
        filter_show_runner.run()

        filter_table = {}
        for filter_param in filter_parser.parse_filter(filter_show_runner.stdout):
            key_item_list = []

            if dataproperty.is_not_empty_string(filter_param.get("network")):
                key_item_list.append("network=" + filter_param.get("network"))

            if dataproperty.is_integer(filter_param.get("port")):
                key_item_list.append(
                    "port={:d}".format(filter_param.get("port")))

            filter_key = ", ".join(key_item_list)
            filter_table[filter_key] = {}
            if filter_param.get("flowid") == qdisc_param.get("parent"):
                work_qdisc_param = dict(qdisc_param)
                del work_qdisc_param["parent"]
                filter_table[filter_key] = work_qdisc_param

        return filter_table

    def __set_pre_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            not dataproperty.is_integer(self.port),
        ]):
            flowid = "{:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id())
        else:
            flowid = "{:x}:2".format(qdisc_major_id)

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 2 u32 match ip {:s} 0.0.0.0/0".format(
                self.__get_network_direction_str()),
            "flowid " + flowid,
        ]

        if all([
            self.direction == TrafficDirection.OUTGOING,
            dataproperty.is_not_empty_string(self.src_network),
        ]):
            command_list.append("match ip src {:s}".format(self.src_network))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_netem(self, qdisc_major_id):
        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
            "handle {:x}:".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id)),
            "netem",
        ]
        if self.packet_loss_rate > 0:
            command_list.append("loss {:f}%".format(self.packet_loss_rate))
        if self.latency_ms > 0:
            command_list.append("delay {:f}ms".format(self.latency_ms))

            if self.latency_distro_ms > 0:
                command_list.append("{:f}ms distribution normal".format(
                    self.latency_distro_ms))

        if self.corruption_rate > 0:
            command_list.append("corrupt {:f}%".format(self.corruption_rate))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_network_filter(self, qdisc_major_id):
        if all([
            dataproperty.is_empty_string(self.network),
            self.port is None,
        ]):
            return 0

        command_list = [
            "tc filter add",
            "dev " + self.__get_tc_device(),
            "protocol ip",
            "parent {:x}:".format(qdisc_major_id),
            "prio 1 u32",
            "flowid {:x}:{:d}".format(
                qdisc_major_id, self.__get_qdisc_minor_id()),
        ]
        if dataproperty.is_not_empty_string(self.network):
            command_list.append("match ip {:s} {:s}".format(
                self.__get_network_direction_str(), self.network))
        if self.port is not None:
            command_list.append("match ip dport {:d} 0xffff".format(self.port))

        return SubprocessRunner(" ".join(command_list)).run()

    def __set_rate(self, qdisc_major_id):
        if dataproperty.is_empty_string(self.bandwidth_rate):
            return 0

        rate_kbps = Humanreadable(kilo_size=1000).humanreadable_to_byte(
            self.bandwidth_rate) / 1000.0
        if rate_kbps <= 0:
            raise ValueError("rate must be greater than zero")

        command_list = [
            "tc qdisc add",
            "dev " + self.__get_tc_device(),
            "parent {:x}:{:d}".format(
                self.__get_netem_qdisc_major_id(qdisc_major_id),
                self.__get_qdisc_minor_id()),
            "handle 20:",
            "tbf",
            "rate {:f}kbit".format(rate_kbps),
            "buffer {:d}".format(
                max(rate_kbps, self.__MIN_BUFFER_BYTE)),  # [byte]
            "limit 10000",
        ]

        return SubprocessRunner(" ".join(command_list)).run()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Generated by Django 1.11.15 on 2020-02-12 09:25
from __future__ import unicode_literals

from django.db import migrations, connection

from bluebottle.clients import properties
from bluebottle.clients.models import Client
from bluebottle.clients.utils import LocalTenant

from django.contrib.auth.models import Permission, Group

def remove_anonymous_permissions(apps, schema_editor):
    permissions = (
        ('assignments', 'api_read_assignment'),
        ('events', 'api_read_event'),
        ('activities', 'api_read_activity'),
        ('funding', 'api_read_funding'),
    )

    tenant = Client.objects.get(schema_name=connection.tenant.schema_name)

    with LocalTenant(tenant):
        if properties.CLOSED_SITE:
            anonymous = Group.objects.get(name='Anonymous')

            for (app, codename) in permissions:
                permission = Permission.objects.get(
                    content_type__app_label=app,
                    codename=codename
                )

                anonymous.permissions.remove(permission)

            anonymous.save()

class Migration(migrations.Migration):

    dependencies = [
        ('activities', '0017_auto_20200205_1054'),
    ]

    operations = [
        migrations.RunPython(remove_anonymous_permissions)
    ]
<commit_msg>Fix migration for new tenants<commit_after># -*- coding: utf-8 -*-
# Generated by Django 1.11.15 on 2020-02-12 09:25
from __future__ import unicode_literals

from django.db import migrations, connection

from bluebottle.clients import properties
from bluebottle.clients.models import Client
from bluebottle.clients.utils import LocalTenant

from django.contrib.auth.models import Permission, Group

def remove_anonymous_permissions(apps, schema_editor):
    permissions = (
        ('assignments', 'api_read_assignment'),
        ('events', 'api_read_event'),
        ('activities', 'api_read_activity'),
        ('funding', 'api_read_funding'),
    )

    tenant = Client.objects.get(schema_name=connection.tenant.schema_name)

    with LocalTenant(tenant):
        if properties.CLOSED_SITE:
            anonymous = Group.objects.get(name='Anonymous')

            for (app, codename) in permissions:
                try:
                    permission = Permission.objects.get(
                        content_type__app_label=app,
                        codename=codename
                    )

                    anonymous.permissions.remove(permission)
                except Permission.DoesNotExist:
                    pass
            anonymous.save()

class Migration(migrations.Migration):

    dependencies = [
        ('activities', '0017_auto_20200205_1054'),
    ]

    operations = [
        migrations.RunPython(remove_anonymous_permissions)
    ]
<|endoftext|>