<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import magic
import os
import mimetypes
import shlex
import subprocess
import time

from kittengroomer import FileBase, KittenGroomerBase, main

UNOCONV = '/usr/bin/unoconv'
LIBREOFFICE = '/usr/bin/libreoffice'
GS = '/usr/bin/gs'
PDF2HTMLEX = '/usr/bin/pdf2htmlEX'
SEVENZ = '/usr/bin/7z'

# Prepare application/<subtype>
mimes_office = ['msword', 'vnd.openxmlformats-officedocument.', 'vnd.ms-',
                'vnd.oasis.opendocument']
mimes_pdf = ['pdf', 'postscript']
mimes_xml = ['xml']
mimes_ms = ['x-dosexec']
mimes_compressed = ['zip', 'x-rar', 'x-bzip2', 'x-lzip', 'x-lzma', 'x-lzop',
                    'x-xz', 'x-compress', 'x-gzip', 'x-tar', 'compressed']
mimes_data = ['octet-stream']

# Aliases
aliases = {
    # Win executables
    'application/x-msdos-program': 'application/x-dosexec',
    'application/x-dosexec': 'application/x-msdos-program'
}

# Sometimes, mimetypes.guess_type is giving unexpected results, such as for the .tar.gz files:
# In [12]: mimetypes.guess_type('toot.tar.gz', strict=False)
# Out[12]: ('application/x-tar', 'gzip')
# It works as expected if you do mimetypes.guess_type('application/gzip', strict=False)
propertype = {'.gz': 'application/gzip'}

class File(FileBase):

    def __init__(self, src_path, dst_path):
        ''' Init file object, set the mimetype '''
        super(File, self).__init__(src_path, dst_path)

        mimetype = magic.from_file(src_path, mime=True).decode("utf-8")
        self.main_type, self.sub_type = mimetype.split('/')
        a, self.extension = os.path.splitext(src_path)
        self.is_recursive = False

        self.log_details.update({'maintype': self.main_type, 'subtype': self.sub_type, 'extension': self.extension})
        # If the mimetype matches as text/*, it will be sent to LibreOffice, no need to cross check the mime/ext
        if self.main_type == 'text':
            return

        # Check correlation known extension => actual mime type
        if propertype.get(self.extension) is not None:
            expected_mimetype = propertype.get(self.extension)
        else:
            expected_mimetype, encoding = mimetypes.guess_type(self.src_path, strict=False)
            if aliases.get(expected_mimetype) is not None:
                expected_mimetype = aliases.get(expected_mimetype)

        is_known_extension = self.extension in mimetypes.types_map.keys()
        if is_known_extension and expected_mimetype != mimetype:
            self.log_details.update({'expected_mimetype': expected_mimetype})
            self.make_dangerous()

        # check correlation actual mime type => known extensions
        if aliases.get(mimetype) is not None:
            mimetype = aliases.get(mimetype)
        expected_extensions = mimetypes.guess_all_extensions(mimetype, strict=False)
        if expected_extensions is not None:
            if len(self.extension) > 0 and self.extension not in expected_extensions:
                self.log_details.update({'expected_extensions': expected_extensions})
                self.make_dangerous()
        else:
            # there are no known extensions associated to this mimetype.
            pass

class KittenGroomer(KittenGroomerBase):

    def __init__(self, root_src=None, root_dst=None, max_recursive=5):
        '''
            Initialize the basics of the conversion process
        '''
        if root_src is None:
            root_src = os.path.join(os.sep, 'media', 'src')
        if root_dst is None:
            root_dst = os.path.join(os.sep, 'media', 'dst')
        super(KittenGroomer, self).__init__(root_src, root_dst)

        self.recursive = 0
        self.max_recursive = max_recursive

        subtypes_apps = [
            (mimes_office, self._office_related),
            (mimes_pdf, self._pdf),
            (mimes_xml, self._office_related),
            (mimes_ms, self._executables),
            (mimes_compressed, self._archive),
            (mimes_data, self._binary_app),
        ]
        self.subtypes_application = self._init_subtypes_application(subtypes_apps)

        self.mime_processing_options = {
            'text': self.text,
            'audio': self.audio,
            'image': self.image,
            'video': self.video,
            'application': self.application,
            'example': self.example,
            'message': self.message,
            'model': self.model,
            'multipart': self.multipart,
            'inode': self.inode,
        }

        unoconv_listener = UNOCONV + ' --listener'
        self._run_process(unoconv_listener, background=True)

    # ##### Helpers #####
    def _init_subtypes_application(self, subtypes_application):
        '''
            Create the Dict to pick the right function based on the sub mime type
        '''
        to_return = {}
        for list_subtypes, fct in subtypes_application:
            for st in list_subtypes:
                to_return[st] = fct
        return to_return

    def _print_log(self):
        '''
            Print the logs related to the current file being processed
        '''
        tmp_log = self.log_name.fields(**self.cur_file.log_details)
        if self.cur_file.log_details.get('dangerous'):
            tmp_log.warning(self.cur_file.log_string)
        elif self.cur_file.log_details.get('unknown') or self.cur_file.log_details.get('binary'):
            tmp_log.info(self.cur_file.log_string)
        else:
            tmp_log.debug(self.cur_file.log_string)

    def _run_process(self, command_line, timeout=0, background=False):
        '''Run subprocess, wait until it finishes'''
        if timeout != 0:
            deadline = time.time() + timeout
        else:
            deadline = None
        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        if background:
            # FIXME: This timer is here to make sure the unoconv listener is properly started.
            time.sleep(10)
            return True
        while True:
            code = p.poll()
            if code is not None:
                break
            if deadline is not None and time.time() > deadline:
                p.kill()
                break
            time.sleep(1)
        return True

    #######################

    # ##### Discarded mime types, reason in the comments ######
    def inode(self):
        ''' Usually empty file. No reason (?) to copy it on the dest key'''
        self.cur_file.log_string += 'Inode file'

    def unknown(self):
        ''' This main type is unknown, that should not happen '''
        self.cur_file.log_string += 'Unknown file'

    # ##### Threated as malicious, no reason to have it on a USB key ######
    def example(self):
        '''Way to process example file'''
        self.cur_file.log_string += 'Example file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    def message(self):
        '''Way to process message file'''
        self.cur_file.log_string += 'Message file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    def model(self):
        '''Way to process model file'''
        self.cur_file.log_string += 'Model file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    def multipart(self):
        '''Way to process multipart file'''
        self.cur_file.log_string += 'Multipart file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    #######################

    # ##### Converted ######
    def text(self):
        ''' LibreOffice should be able to open all the files '''
        self.cur_file.log_string += 'Text file'
        self._office_related()

    def application(self):
        ''' Everything can be there, using the subtype to decide '''
        for subtype, fct in list(self.subtypes_application.items()):
            if subtype in self.cur_file.sub_type:
                fct()
                self.cur_file.log_string += 'Application file'
                return
        self.cur_file.log_string += 'Unknown Application file'
        self._unknown_app()

    def _executables(self):
        '''Way to process executable file'''
        self.cur_file.add_log_details('processing_type', 'executable')
        self.cur_file.make_dangerous()
        self._safe_copy()

    def _office_related(self):
        '''Way to process all the files LibreOffice can handle'''
        self.cur_file.add_log_details('processing_type', 'office')
        dst_dir, filename = os.path.split(self.cur_file.dst_path)
        tmpdir = os.path.join(dst_dir, 'temp')
        name, ext = os.path.splitext(filename)
        tmppath = os.path.join(tmpdir, name + '.pdf')
        self._safe_mkdir(tmpdir)
        lo_command = '{} --format pdf -eSelectPdfVersion=1 --output {} {}'.format(
            UNOCONV, tmppath, self.cur_file.src_path)
        self._run_process(lo_command)
        self._pdfa(tmppath)
        self._safe_rmtree(tmpdir)

    def _pdfa(self, tmpsrcpath):
        '''Way to process PDF/A file'''
        pdf_command = '{} --dest-dir / {} {}'.format(PDF2HTMLEX, tmpsrcpath,
                                                     self.cur_file.dst_path + '.html')
        self._run_process(pdf_command)

    def _pdf(self):
        '''Way to process PDF file'''
        self.cur_file.add_log_details('processing_type', 'pdf')
        dst_dir, filename = os.path.split(self.cur_file.dst_path)
        tmpdir = os.path.join(dst_dir, 'temp')
        tmppath = os.path.join(tmpdir, filename)
        self._safe_mkdir(tmpdir)
        # The magic comes from here: http://svn.ghostscript.com/ghostscript/trunk/gs/doc/Ps2pdf.htm#PDFA
        curdir = os.getcwd()
        os.chdir(self.ressources_path)
        gs_command = '{} -dPDFA -dBATCH -dNOPAUSE -dNOOUTERSAVE -sProcessColorModel=DeviceCMYK -sDEVICE=pdfwrite -sPDFACompatibilityPolicy=1 -sOutputFile={} ./PDFA_def.ps {}'.format(
            GS, os.path.join(curdir, tmppath), os.path.join(curdir, self.cur_file.src_path))
        self._run_process(gs_command)
        os.chdir(curdir)
        self._pdfa(tmppath)
        self._safe_rmtree(tmpdir)

    def _archive(self):
        '''Way to process Archive'''
        self.cur_file.add_log_details('processing_type', 'archive')
        self.cur_file.is_recursive = True
        self.cur_file.log_string += 'Archive extracted, processing content.'
        tmpdir = self.cur_file.dst_path + '_temp'
        self._safe_mkdir(tmpdir)
        extract_command = '{} -p1 x {} -o{} -bd -aoa'.format(SEVENZ, self.cur_file.src_path, tmpdir)
        self._run_process(extract_command)
        self.recursive += 1
        self.processdir(tmpdir, self.cur_file.dst_path)
        self.recursive -= 1
        self._safe_rmtree(tmpdir)

    def _unknown_app(self):
        '''Way to process an unknown file'''
        self.cur_file.make_unknown()
        self._safe_copy()

    def _binary_app(self):
        '''Way to process an unknown binary file'''
        self.cur_file.make_binary()
        self._safe_copy()

    #######################

    # ##### Not converted, checking the mime type ######
    def audio(self):
        '''Way to process an audio file'''
        self.cur_file.log_string += 'Audio file'
        self._media_processing()

    def image(self):
        '''Way to process an image'''
        self.cur_file.log_string += 'Image file'
        self._media_processing()

    def video(self):
        '''Way to process a video'''
        self.cur_file.log_string += 'Video file'
        self._media_processing()

    def _media_processing(self):
        '''Generic way to process all the media files'''
        self.cur_file.add_log_details('processing_type', 'media')
        self._safe_copy()

    #######################

    def processdir(self, src_dir=None, dst_dir=None):
        '''
            Main function doing the processing
        '''
        if src_dir is None:
            src_dir = self.src_root_dir
        if dst_dir is None:
            dst_dir = self.dst_root_dir

        if self.recursive > 0:
            self._print_log()

        if self.recursive >= self.max_recursive:
            self.cur_log.warning('ARCHIVE BOMB.')
            self.cur_log.warning('The content of the archive contains recursively other archives.')
            self.cur_log.warning('This is a bad sign so the archive is not extracted to the destination key.')
            self._safe_rmtree(src_dir)
            if src_dir.endswith('_temp'):
                archbomb_path = src_dir[:-len('_temp')]
                self._safe_remove(archbomb_path)

        for srcpath in self._list_all_files(src_dir):
            self.cur_file = File(srcpath, srcpath.replace(src_dir, dst_dir))

            self.log_name.info('Processing {} ({}/{})', srcpath.replace(src_dir + '/', ''),
                               self.cur_file.main_type, self.cur_file.sub_type)
            if self.cur_file.log_details.get('dangerous') is None:
                self.mime_processing_options.get(self.cur_file.main_type, self.unknown)()
            else:
                self._safe_copy()
            if not self.cur_file.is_recursive:
                self._print_log()

if __name__ == '__main__':
    main(KittenGroomer, 'Generic version of the KittenGroomer. Convert and rename files.')
<commit_msg>Make GS conversion safer<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import magic
import os
import mimetypes
import shlex
import subprocess
import time

from kittengroomer import FileBase, KittenGroomerBase, main

UNOCONV = '/usr/bin/unoconv'
LIBREOFFICE = '/usr/bin/libreoffice'
GS = '/usr/bin/gs'
PDF2HTMLEX = '/usr/bin/pdf2htmlEX'
SEVENZ = '/usr/bin/7z'

# Prepare application/<subtype>
mimes_office = ['msword', 'vnd.openxmlformats-officedocument.', 'vnd.ms-',
                'vnd.oasis.opendocument']
mimes_pdf = ['pdf', 'postscript']
mimes_xml = ['xml']
mimes_ms = ['x-dosexec']
mimes_compressed = ['zip', 'x-rar', 'x-bzip2', 'x-lzip', 'x-lzma', 'x-lzop',
                    'x-xz', 'x-compress', 'x-gzip', 'x-tar', 'compressed']
mimes_data = ['octet-stream']

# Aliases
aliases = {
    # Win executables
    'application/x-msdos-program': 'application/x-dosexec',
    'application/x-dosexec': 'application/x-msdos-program'
}

# Sometimes, mimetypes.guess_type is giving unexpected results, such as for the .tar.gz files:
# In [12]: mimetypes.guess_type('toot.tar.gz', strict=False)
# Out[12]: ('application/x-tar', 'gzip')
# It works as expected if you do mimetypes.guess_type('application/gzip', strict=False)
propertype = {'.gz': 'application/gzip'}

class File(FileBase):

    def __init__(self, src_path, dst_path):
        ''' Init file object, set the mimetype '''
        super(File, self).__init__(src_path, dst_path)

        mimetype = magic.from_file(src_path, mime=True).decode("utf-8")
        self.main_type, self.sub_type = mimetype.split('/')
        a, self.extension = os.path.splitext(src_path)
        self.is_recursive = False

        self.log_details.update({'maintype': self.main_type, 'subtype': self.sub_type, 'extension': self.extension})
        # If the mimetype matches as text/*, it will be sent to LibreOffice, no need to cross check the mime/ext
        if self.main_type == 'text':
            return

        # Check correlation known extension => actual mime type
        if propertype.get(self.extension) is not None:
            expected_mimetype = propertype.get(self.extension)
        else:
            expected_mimetype, encoding = mimetypes.guess_type(self.src_path, strict=False)
            if aliases.get(expected_mimetype) is not None:
                expected_mimetype = aliases.get(expected_mimetype)

        is_known_extension = self.extension in mimetypes.types_map.keys()
        if is_known_extension and expected_mimetype != mimetype:
            self.log_details.update({'expected_mimetype': expected_mimetype})
            self.make_dangerous()

        # check correlation actual mime type => known extensions
        if aliases.get(mimetype) is not None:
            mimetype = aliases.get(mimetype)
        expected_extensions = mimetypes.guess_all_extensions(mimetype, strict=False)
        if expected_extensions is not None:
            if len(self.extension) > 0 and self.extension not in expected_extensions:
                self.log_details.update({'expected_extensions': expected_extensions})
                self.make_dangerous()
        else:
            # there are no known extensions associated to this mimetype.
            pass

class KittenGroomer(KittenGroomerBase):

    def __init__(self, root_src=None, root_dst=None, max_recursive=5):
        '''
            Initialize the basics of the conversion process
        '''
        if root_src is None:
            root_src = os.path.join(os.sep, 'media', 'src')
        if root_dst is None:
            root_dst = os.path.join(os.sep, 'media', 'dst')
        super(KittenGroomer, self).__init__(root_src, root_dst)

        self.recursive = 0
        self.max_recursive = max_recursive

        subtypes_apps = [
            (mimes_office, self._office_related),
            (mimes_pdf, self._pdf),
            (mimes_xml, self._office_related),
            (mimes_ms, self._executables),
            (mimes_compressed, self._archive),
            (mimes_data, self._binary_app),
        ]
        self.subtypes_application = self._init_subtypes_application(subtypes_apps)

        self.mime_processing_options = {
            'text': self.text,
            'audio': self.audio,
            'image': self.image,
            'video': self.video,
            'application': self.application,
            'example': self.example,
            'message': self.message,
            'model': self.model,
            'multipart': self.multipart,
            'inode': self.inode,
        }

        unoconv_listener = UNOCONV + ' --listener'
        self._run_process(unoconv_listener, background=True)

    # ##### Helpers #####
    def _init_subtypes_application(self, subtypes_application):
        '''
            Create the Dict to pick the right function based on the sub mime type
        '''
        to_return = {}
        for list_subtypes, fct in subtypes_application:
            for st in list_subtypes:
                to_return[st] = fct
        return to_return

    def _print_log(self):
        '''
            Print the logs related to the current file being processed
        '''
        tmp_log = self.log_name.fields(**self.cur_file.log_details)
        if self.cur_file.log_details.get('dangerous'):
            tmp_log.warning(self.cur_file.log_string)
        elif self.cur_file.log_details.get('unknown') or self.cur_file.log_details.get('binary'):
            tmp_log.info(self.cur_file.log_string)
        else:
            tmp_log.debug(self.cur_file.log_string)

    def _run_process(self, command_line, timeout=0, background=False):
        '''Run subprocess, wait until it finishes'''
        if timeout != 0:
            deadline = time.time() + timeout
        else:
            deadline = None
        args = shlex.split(command_line)
        p = subprocess.Popen(args)
        if background:
            # FIXME: This timer is here to make sure the unoconv listener is properly started.
            time.sleep(10)
            return True
        while True:
            code = p.poll()
            if code is not None:
                break
            if deadline is not None and time.time() > deadline:
                p.kill()
                break
            time.sleep(1)
        return True

    #######################

    # ##### Discarded mime types, reason in the comments ######
    def inode(self):
        ''' Usually empty file. No reason (?) to copy it on the dest key'''
        self.cur_file.log_string += 'Inode file'

    def unknown(self):
        ''' This main type is unknown, that should not happen '''
        self.cur_file.log_string += 'Unknown file'

    # ##### Threated as malicious, no reason to have it on a USB key ######
    def example(self):
        '''Way to process example file'''
        self.cur_file.log_string += 'Example file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    def message(self):
        '''Way to process message file'''
        self.cur_file.log_string += 'Message file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    def model(self):
        '''Way to process model file'''
        self.cur_file.log_string += 'Model file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    def multipart(self):
        '''Way to process multipart file'''
        self.cur_file.log_string += 'Multipart file'
        self.cur_file.make_dangerous()
        self._safe_copy()

    #######################

    # ##### Converted ######
    def text(self):
        ''' LibreOffice should be able to open all the files '''
        self.cur_file.log_string += 'Text file'
        self._office_related()

    def application(self):
        ''' Everything can be there, using the subtype to decide '''
        for subtype, fct in list(self.subtypes_application.items()):
            if subtype in self.cur_file.sub_type:
                fct()
                self.cur_file.log_string += 'Application file'
                return
        self.cur_file.log_string += 'Unknown Application file'
        self._unknown_app()

    def _executables(self):
        '''Way to process executable file'''
        self.cur_file.add_log_details('processing_type', 'executable')
        self.cur_file.make_dangerous()
        self._safe_copy()

    def _office_related(self):
        '''Way to process all the files LibreOffice can handle'''
        self.cur_file.add_log_details('processing_type', 'office')
        dst_dir, filename = os.path.split(self.cur_file.dst_path)
        tmpdir = os.path.join(dst_dir, 'temp')
        name, ext = os.path.splitext(filename)
        tmppath = os.path.join(tmpdir, name + '.pdf')
        self._safe_mkdir(tmpdir)
        lo_command = '{} --format pdf -eSelectPdfVersion=1 --output {} {}'.format(
            UNOCONV, tmppath, self.cur_file.src_path)
        self._run_process(lo_command)
        self._pdfa(tmppath)
        self._safe_rmtree(tmpdir)

    def _pdfa(self, tmpsrcpath):
        '''Way to process PDF/A file'''
        pdf_command = '{} --dest-dir / {} {}'.format(PDF2HTMLEX, tmpsrcpath,
                                                     self.cur_file.dst_path + '.html')
        self._run_process(pdf_command)

    def _pdf(self):
        '''Way to process PDF file'''
        self.cur_file.add_log_details('processing_type', 'pdf')
        dst_dir, filename = os.path.split(self.cur_file.dst_path)
        tmpdir = os.path.join(dst_dir, 'temp')
        tmppath = os.path.join(tmpdir, filename)
        self._safe_mkdir(tmpdir)
        # The magic comes from here: http://svn.ghostscript.com/ghostscript/trunk/gs/doc/Ps2pdf.htm#PDFA
        curdir = os.getcwd()
        os.chdir(self.ressources_path)
        gs_command = '{} -dPDFA -dQUIET -dSAFER -dBATCH -dNOPAUSE -dNOOUTERSAVE -sProcessColorModel=DeviceCMYK -sDEVICE=pdfwrite -sPDFACompatibilityPolicy=1 -sOutputFile={} ./PDFA_def.ps {}'.format(
            GS, os.path.join(curdir, tmppath), os.path.join(curdir, self.cur_file.src_path))
        self._run_process(gs_command)
        os.chdir(curdir)
        self._pdfa(tmppath)
        self._safe_rmtree(tmpdir)

    def _archive(self):
        '''Way to process Archive'''
        self.cur_file.add_log_details('processing_type', 'archive')
        self.cur_file.is_recursive = True
        self.cur_file.log_string += 'Archive extracted, processing content.'
        tmpdir = self.cur_file.dst_path + '_temp'
        self._safe_mkdir(tmpdir)
        extract_command = '{} -p1 x {} -o{} -bd -aoa'.format(SEVENZ, self.cur_file.src_path, tmpdir)
        self._run_process(extract_command)
        self.recursive += 1
        self.processdir(tmpdir, self.cur_file.dst_path)
        self.recursive -= 1
        self._safe_rmtree(tmpdir)

    def _unknown_app(self):
        '''Way to process an unknown file'''
        self.cur_file.make_unknown()
        self._safe_copy()

    def _binary_app(self):
        '''Way to process an unknown binary file'''
        self.cur_file.make_binary()
        self._safe_copy()

    #######################

    # ##### Not converted, checking the mime type ######
    def audio(self):
        '''Way to process an audio file'''
        self.cur_file.log_string += 'Audio file'
        self._media_processing()

    def image(self):
        '''Way to process an image'''
        self.cur_file.log_string += 'Image file'
        self._media_processing()

    def video(self):
        '''Way to process a video'''
        self.cur_file.log_string += 'Video file'
        self._media_processing()

    def _media_processing(self):
        '''Generic way to process all the media files'''
        self.cur_file.add_log_details('processing_type', 'media')
        self._safe_copy()

    #######################

    def processdir(self, src_dir=None, dst_dir=None):
        '''
            Main function doing the processing
        '''
        if src_dir is None:
            src_dir = self.src_root_dir
        if dst_dir is None:
            dst_dir = self.dst_root_dir

        if self.recursive > 0:
            self._print_log()

        if self.recursive >= self.max_recursive:
            self.cur_log.warning('ARCHIVE BOMB.')
            self.cur_log.warning('The content of the archive contains recursively other archives.')
            self.cur_log.warning('This is a bad sign so the archive is not extracted to the destination key.')
            self._safe_rmtree(src_dir)
            if src_dir.endswith('_temp'):
                archbomb_path = src_dir[:-len('_temp')]
                self._safe_remove(archbomb_path)

        for srcpath in self._list_all_files(src_dir):
            self.cur_file = File(srcpath, srcpath.replace(src_dir, dst_dir))

            self.log_name.info('Processing {} ({}/{})', srcpath.replace(src_dir + '/', ''),
                               self.cur_file.main_type, self.cur_file.sub_type)
            if self.cur_file.log_details.get('dangerous') is None:
                self.mime_processing_options.get(self.cur_file.main_type, self.unknown)()
            else:
                self._safe_copy()
            if not self.cur_file.is_recursive:
                self._print_log()

if __name__ == '__main__':
    main(KittenGroomer, 'Generic version of the KittenGroomer. Convert and rename files.')
<|endoftext|>


<commit_before>#!/bin/env python3

#
# Moses: A simple Socks5 proxy that encrypts your connections using TLS.
#
# Use the `-h` option to see usage info.
#

#
#   Then Moses stretched out his hand over the sea, and all that night
#   the Lord drove the sea back with a strong east wind and turned it
#   into dry land. The waters were divided, and the Israelites went
#   through the sea on dry ground, with a wall of water on their right
#   and on their left. ( Exodus 14:21-22 )
#

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

import asyncio
from concurrent.futures import FIRST_COMPLETED
import ssl
import functools
import logging
import traceback
import argparse
import sys

DEFAULT_BLOCK_SIZE = 2048

logger = logging.getLogger(__file__)

def socks_check_version(ver):
    if ver != 0x05:
        logger.debug('Protocol version %d not supported', ver)
        raise VersionNotSupportedError('Bad version: {}'.format(ver))

@asyncio.coroutine
def socks_recv_auth_method_list(reader):
    header = yield from reader.readexactly(2)
    socks_check_version(header[0])
    methods = header[1]
    if methods > 0:
        method_list = yield from reader.readexactly(methods)
    else:
        method_list = b''
    return method_list

@asyncio.coroutine
def socks_recv_request(reader):
    header = yield from reader.readexactly(4)
    socks_check_version(header[0])

    if header[3] == 0x01:       # ip
        ip_binary = yield from reader.readexactly(4)
        address = '{}.{}.{}.{}'.format(
            ip_binary[0], ip_binary[1], ip_binary[2], ip_binary[3])
    elif header[3] == 0x03:     # domain name
        domain_len = (yield from reader.readexactly(1))[0]
        address = (yield from reader.readexactly(domain_len)).decode()
    elif header[3] == 0x04:     # ipv6
        ip_binary = yield from reader.readexactly(16)
        address = ('{:02x}{:02x}:' * 7 + '{:02x}{:02x}').format(
            ip_binary[0], ip_binary[1], ip_binary[2], ip_binary[3],
            ip_binary[4], ip_binary[5], ip_binary[6], ip_binary[7],
            ip_binary[8], ip_binary[9], ip_binary[10], ip_binary[11],
            ip_binary[12], ip_binary[13], ip_binary[14], ip_binary[15])

    port_binary = yield from reader.readexactly(2)
    port = (port_binary[0] << 8) + port_binary[1]

    #       command,   addr_type
    return (header[1], header[3], address, port)

@asyncio.coroutine
def socks_handshake(reader, writer):
    try:
        method_list = yield from socks_recv_auth_method_list(reader)
    except:
        logger.debug(traceback.format_exc())
        return None

    if 0x0 not in method_list:
        logger.debug('No proper method found')
        yield from sync_write(writer, b'\x05\xff')
        return None

    yield from sync_write(writer, b'\x05\x00')

    try:
        command, addr_type, address, port = \
                yield from socks_recv_request(reader)
    except:
        logger.debug(traceback.format_exc())
        return None

    if command != 0x01:     # XXX: only support the CONNECT command
        logger.debug('Bad command: %d', command)
        yield from sync_write(
            writer, b'\x05\x07\x00\x01\x00\x00\x00\x00\x00\x00')
        return None

    return (command, addr_type, address, port)

@asyncio.coroutine
def socks_handshake_done(socks_req, reader, writer):
    yield from sync_write(
        writer, b'\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00')

@asyncio.coroutine
def sync_write(writer, data):
    writer.write(data)
    return (yield from writer.drain())

@asyncio.coroutine
def streaming(reader, writer, block_size):
    data = yield from reader.read(block_size)
    while len(data) > 0:
        yield from sync_write(writer, data)
        data = yield from reader.read(block_size)
    if writer.can_write_eof():
        writer.write_eof()
        return True
    else:
        return False

@asyncio.coroutine
def do_connect(addr, port, ssl=None):
    connect_op = asyncio.async(asyncio.open_connection(addr, port, ssl=ssl))
    try:
        proxy_reader, proxy_writer = \
                yield from asyncio.wait_for(connect_op, None)
    except:
        connect_op.cancel()     # for python version < 3.4.3
        logger.debug(traceback.format_exc())
        return None

    if ssl is not None:
        cur_cipher = proxy_writer.get_extra_info('cipher')
        logger.debug('Current cipher: %s', cur_cipher)

    return (proxy_reader, proxy_writer)

@asyncio.coroutine
def do_streaming(reader, writer, remote_reader, remote_writer, bs):
    stream_up = asyncio.async(streaming(reader, remote_writer, bs))
    stream_down = asyncio.async(streaming(remote_reader, writer, bs))

    done, pending = yield from \
            asyncio.wait([stream_up, stream_down], return_when=FIRST_COMPLETED)

    for f in done:
        try:
            f.result()
        except:
            logger.debug(traceback.format_exc())

    for f in pending:
        f.cancel()

@asyncio.coroutine
def server_connection_cb(reader, writer, bs=DEFAULT_BLOCK_SIZE):
    socks_req = yield from socks_handshake(reader, writer)
    if socks_req is None:
        writer.close()
        return

    command, addr_type, address, port = socks_req

    logger.debug('Connecting to %s:%d', address, port)

    remote_rw = yield from do_connect(address, port)
    if remote_rw is None:
        writer.close()
        return

    yield from socks_handshake_done(socks_req, reader, writer)

    yield from do_streaming(reader, writer, remote_rw[0], remote_rw[1], bs)

    remote_rw[1].close()
    writer.close()

@asyncio.coroutine
def client_connection_cb(reader, writer, server_addr, bs=DEFAULT_BLOCK_SIZE, ssl_ctx=None):
    remote_rw = yield from do_connect(server_addr[0], server_addr[1], ssl=ssl_ctx)
    if remote_rw is None:
        writer.close()
        return

    yield from do_streaming(reader, writer, remote_rw[0], remote_rw[1], bs)

    remote_rw[1].close()
    writer.close()

def build_ssl_ctx(my_certs_file, peer_certs_file, ciphers=None):
    ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
    if ciphers is not None:
        ssl_ctx.set_ciphers(ciphers)
    ssl_ctx.options |= ssl.OP_NO_SSLv2
    ssl_ctx.options |= ssl.OP_NO_SSLv3
    ssl_ctx.verify_mode = ssl.CERT_REQUIRED
    ssl_ctx.check_hostname = False
    ssl_ctx.load_cert_chain(my_certs_file)
    ssl_ctx.load_verify_locations(peer_certs_file)
    return ssl_ctx

def server_main(loop, args):
    logger.info('Moses server listening at %s', args.bind)

    if args.forward is not None:
        logger.info('Forwarding client connections to %s', args.forward)
        try:
            forward_addr = parse_ip_port(args.forward)
        except:
            logger.error('Bad forwarding address: %s', args.forward)
            sys.exit(1)
        cb = functools.partial(client_connection_cb,
                server_addr=forward_addr, bs=args.block_size)
    else:
        cb = functools.partial(server_connection_cb, bs=args.block_size)

    if args.no_tls:
        logger.warning('Connections from clients are NOT encrypted')
        ssl_ctx = None
    else:
        ssl_ctx = build_ssl_ctx(args.local_cert, args.remote_cert, args.ciphers)

    local_ip, local_port = parse_ip_port(args.bind)
    starter = asyncio.start_server(cb, local_ip, local_port,
            ssl=ssl_ctx,
            backlog=args.backlog,
            reuse_address=True,
            loop=loop)
    return loop.run_until_complete(starter)

def client_main(loop, args):
    logger.info('Moses client listening at %s', args.bind)
    logger.info('Forwarding to %s', args.peer)

    if args.peer is None:
        logger.error('No peer (server) address provided')
        sys.exit(1)

    try:
        server_addr = parse_ip_port(args.peer)
    except:
        logger.error('Bad peer address: %s', args.peer)
        sys.exit(1)

    if args.no_tls:
        logger.warning('Connections to the server are NOT encrypted')
        ssl_ctx = None
    else:
        ssl_ctx = build_ssl_ctx(args.local_cert, args.remote_cert, args.ciphers)

    cb = functools.partial(client_connection_cb,
            server_addr=server_addr,
            bs=args.block_size,
            ssl_ctx=ssl_ctx)

    local_ip, local_port = parse_ip_port(args.bind)
    starter = asyncio.start_server(cb, local_ip, local_port,
            backlog=args.backlog,
            reuse_address=True,
            loop=loop)
    return loop.run_until_complete(starter)

def parse_arguments():
    parser = argparse.ArgumentParser()

    common_group = parser.add_argument_group('Common Options')

    mode_group = common_group.add_mutually_exclusive_group(required=True)
    mode_group.add_argument('-c', '--client',
            help='Client mode',
            action='store_true')
    mode_group.add_argument('-s', '--server',
            help='Server mode',
            action='store_true')

    common_group.add_argument('-b', '--bind',
            metavar='<ADDRESS>:<PORT>',
            help='IP & port to bind (default: <all interfaces>:1080)',
            default=':1080',
            type=str)
    common_group.add_argument('-n', '--no-tls',
            help='Do not use TLS encryption',
            action='store_true')
    common_group.add_argument('-l', '--local-cert',
            help='Local SSL certificates (default: ./local.pem)',
            default='./local.pem',
            type=str)
    common_group.add_argument('-r', '--remote-cert',
            help='Remote SSL certificates (default: ./remote.pem)',
            default='./remote.pem',
            type=str)
    common_group.add_argument('-e', '--ciphers',
            help='Ciphers to use for encryption. '
                 'Run `openssl ciphers` to see available ciphers',
            type=str)
    common_group.add_argument('--backlog',
            help='Backlog for the listening socket (default: 128)',
            default=128,
            type=int)
    common_group.add_argument('--loglevel',
            help='Log level (default: info)',
            default='info',
            type=str,
            choices=[
                'critical', 'fatal', 'error',
                'warning', 'info', 'debug',
                ])
    common_group.add_argument('--block-size',
            help='Block size for data streaming, in bytes (default: 2048)',
            default=DEFAULT_BLOCK_SIZE,
            type=int)

    client_group = parser.add_argument_group('Client Options')

    client_group.add_argument('-p', '--peer',
            metavar='<ADDRESS>:<PORT>',
            help='Peer (server) address',
            type=str)

    server_group = parser.add_argument_group('Server Options')

    server_group.add_argument('-f', '--forward',
            metavar='<ADDRESS>:<PORT>',
            help='Simply forward all connections to the given address',
            type=str)

    return parser.parse_args()

def parse_ip_port(addr_str):
    last_col = addr_str.rfind(':')
    return (addr_str[0:last_col], int(addr_str[last_col+1:]))

def main():
    args = parse_arguments()
    logging.basicConfig(
            level=args.loglevel.upper(),
            format='%(asctime)s:%(levelname)s:%(name)s:%(message)s')

    loop = asyncio.get_event_loop()

    if args.client:
        server = client_main(loop, args)
    elif args.server:
        server = server_main(loop, args)

    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass

    server.close()
    loop.run_until_complete(server.wait_closed())
    loop.close()

if __name__ == '__main__':
    main()
<commit_msg>Use a more 'portable' location of env command<commit_after>#!/usr/bin/env python3

#
# Moses: A simple Socks5 proxy that encrypts your connections using TLS.
#
# Use the `-h` option to see usage info.
#

#
#   Then Moses stretched out his hand over the sea, and all that night
#   the Lord drove the sea back with a strong east wind and turned it
#   into dry land. The waters were divided, and the Israelites went
#   through the sea on dry ground, with a wall of water on their right
#   and on their left. ( Exodus 14:21-22 )
#

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

import asyncio
from concurrent.futures import FIRST_COMPLETED
import ssl
import functools
import logging
import traceback
import argparse
import sys

DEFAULT_BLOCK_SIZE = 2048

logger = logging.getLogger(__file__)

def socks_check_version(ver):
    if ver != 0x05:
        logger.debug('Protocol version %d not supported', ver)
        raise VersionNotSupportedError('Bad version: {}'.format(ver))

@asyncio.coroutine
def socks_recv_auth_method_list(reader):
    header = yield from reader.readexactly(2)
    socks_check_version(header[0])
    methods = header[1]
    if methods > 0:
        method_list = yield from reader.readexactly(methods)
    else:
        method_list = b''
    return method_list

@asyncio.coroutine
def socks_recv_request(reader):
    header = yield from reader.readexactly(4)
    socks_check_version(header[0])

    if header[3] == 0x01:       # ip
        ip_binary = yield from reader.readexactly(4)
        address = '{}.{}.{}.{}'.format(
            ip_binary[0], ip_binary[1], ip_binary[2], ip_binary[3])
    elif header[3] == 0x03:     # domain name
        domain_len = (yield from reader.readexactly(1))[0]
        address = (yield from reader.readexactly(domain_len)).decode()
    elif header[3] == 0x04:     # ipv6
        ip_binary = yield from reader.readexactly(16)
        address = ('{:02x}{:02x}:' * 7 + '{:02x}{:02x}').format(
            ip_binary[0], ip_binary[1], ip_binary[2], ip_binary[3],
            ip_binary[4], ip_binary[5], ip_binary[6], ip_binary[7],
            ip_binary[8], ip_binary[9], ip_binary[10], ip_binary[11],
            ip_binary[12], ip_binary[13], ip_binary[14], ip_binary[15])

    port_binary = yield from reader.readexactly(2)
    port = (port_binary[0] << 8) + port_binary[1]

    #       command,   addr_type
    return (header[1], header[3], address, port)

@asyncio.coroutine
def socks_handshake(reader, writer):
    try:
        method_list = yield from socks_recv_auth_method_list(reader)
    except:
        logger.debug(traceback.format_exc())
        return None

    if 0x0 not in method_list:
        logger.debug('No proper method found')
        yield from sync_write(writer, b'\x05\xff')
        return None

    yield from sync_write(writer, b'\x05\x00')

    try:
        command, addr_type, address, port = \
                yield from socks_recv_request(reader)
    except:
        logger.debug(traceback.format_exc())
        return None

    if command != 0x01:     # XXX: only support the CONNECT command
        logger.debug('Bad command: %d', command)
        yield from sync_write(
            writer, b'\x05\x07\x00\x01\x00\x00\x00\x00\x00\x00')
        return None

    return (command, addr_type, address, port)

@asyncio.coroutine
def socks_handshake_done(socks_req, reader, writer):
    yield from sync_write(
        writer, b'\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00')

@asyncio.coroutine
def sync_write(writer, data):
    writer.write(data)
    return (yield from writer.drain())

@asyncio.coroutine
def streaming(reader, writer, block_size):
    data = yield from reader.read(block_size)
    while len(data) > 0:
        yield from sync_write(writer, data)
        data = yield from reader.read(block_size)
    if writer.can_write_eof():
        writer.write_eof()
        return True
    else:
        return False

@asyncio.coroutine
def do_connect(addr, port, ssl=None):
    connect_op = asyncio.async(asyncio.open_connection(addr, port, ssl=ssl))
    try:
        proxy_reader, proxy_writer = \
                yield from asyncio.wait_for(connect_op, None)
    except:
        connect_op.cancel()     # for python version < 3.4.3
        logger.debug(traceback.format_exc())
        return None

    if ssl is not None:
        cur_cipher = proxy_writer.get_extra_info('cipher')
        logger.debug('Current cipher: %s', cur_cipher)

    return (proxy_reader, proxy_writer)

@asyncio.coroutine
def do_streaming(reader, writer, remote_reader, remote_writer, bs):
    stream_up = asyncio.async(streaming(reader, remote_writer, bs))
    stream_down = asyncio.async(streaming(remote_reader, writer, bs))

    done, pending = yield from \
            asyncio.wait([stream_up, stream_down], return_when=FIRST_COMPLETED)

    for f in done:
        try:
            f.result()
        except:
            logger.debug(traceback.format_exc())

    for f in pending:
        f.cancel()

@asyncio.coroutine
def server_connection_cb(reader, writer, bs=DEFAULT_BLOCK_SIZE):
    socks_req = yield from socks_handshake(reader, writer)
    if socks_req is None:
        writer.close()
        return

    command, addr_type, address, port = socks_req

    logger.debug('Connecting to %s:%d', address, port)

    remote_rw = yield from do_connect(address, port)
    if remote_rw is None:
        writer.close()
        return

    yield from socks_handshake_done(socks_req, reader, writer)

    yield from do_streaming(reader, writer, remote_rw[0], remote_rw[1], bs)

    remote_rw[1].close()
    writer.close()

@asyncio.coroutine
def client_connection_cb(reader, writer, server_addr, bs=DEFAULT_BLOCK_SIZE, ssl_ctx=None):
    remote_rw = yield from do_connect(server_addr[0], server_addr[1], ssl=ssl_ctx)
    if remote_rw is None:
        writer.close()
        return

    yield from do_streaming(reader, writer, remote_rw[0], remote_rw[1], bs)

    remote_rw[1].close()
    writer.close()

def build_ssl_ctx(my_certs_file, peer_certs_file, ciphers=None):
    ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
    if ciphers is not None:
        ssl_ctx.set_ciphers(ciphers)
    ssl_ctx.options |= ssl.OP_NO_SSLv2
    ssl_ctx.options |= ssl.OP_NO_SSLv3
    ssl_ctx.verify_mode = ssl.CERT_REQUIRED
    ssl_ctx.check_hostname = False
    ssl_ctx.load_cert_chain(my_certs_file)
    ssl_ctx.load_verify_locations(peer_certs_file)
    return ssl_ctx

def server_main(loop, args):
    logger.info('Moses server listening at %s', args.bind)

    if args.forward is not None:
        logger.info('Forwarding client connections to %s', args.forward)
        try:
            forward_addr = parse_ip_port(args.forward)
        except:
            logger.error('Bad forwarding address: %s', args.forward)
            sys.exit(1)
        cb = functools.partial(client_connection_cb,
                server_addr=forward_addr, bs=args.block_size)
    else:
        cb = functools.partial(server_connection_cb, bs=args.block_size)

    if args.no_tls:
        logger.warning('Connections from clients are NOT encrypted')
        ssl_ctx = None
    else:
        ssl_ctx = build_ssl_ctx(args.local_cert, args.remote_cert, args.ciphers)

    local_ip, local_port = parse_ip_port(args.bind)
    starter = asyncio.start_server(cb, local_ip, local_port,
            ssl=ssl_ctx,
            backlog=args.backlog,
            reuse_address=True,
            loop=loop)
    return loop.run_until_complete(starter)

def client_main(loop, args):
    logger.info('Moses client listening at %s', args.bind)
    logger.info('Forwarding to %s', args.peer)

    if args.peer is None:
        logger.error('No peer (server) address provided')
        sys.exit(1)

    try:
        server_addr = parse_ip_port(args.peer)
    except:
        logger.error('Bad peer address: %s', args.peer)
        sys.exit(1)

    if args.no_tls:
        logger.warning('Connections to the server are NOT encrypted')
        ssl_ctx = None
    else:
        ssl_ctx = build_ssl_ctx(args.local_cert, args.remote_cert, args.ciphers)

    cb = functools.partial(client_connection_cb,
            server_addr=server_addr,
            bs=args.block_size,
            ssl_ctx=ssl_ctx)

    local_ip, local_port = parse_ip_port(args.bind)
    starter = asyncio.start_server(cb, local_ip, local_port,
            backlog=args.backlog,
            reuse_address=True,
            loop=loop)
    return loop.run_until_complete(starter)

def parse_arguments():
    parser = argparse.ArgumentParser()

    common_group = parser.add_argument_group('Common Options')

    mode_group = common_group.add_mutually_exclusive_group(required=True)
    mode_group.add_argument('-c', '--client',
            help='Client mode',
            action='store_true')
    mode_group.add_argument('-s', '--server',
            help='Server mode',
            action='store_true')

    common_group.add_argument('-b', '--bind',
            metavar='<ADDRESS>:<PORT>',
            help='IP & port to bind (default: <all interfaces>:1080)',
            default=':1080',
            type=str)
    common_group.add_argument('-n', '--no-tls',
            help='Do not use TLS encryption',
            action='store_true')
    common_group.add_argument('-l', '--local-cert',
            help='Local SSL certificates (default: ./local.pem)',
            default='./local.pem',
            type=str)
    common_group.add_argument('-r', '--remote-cert',
            help='Remote SSL certificates (default: ./remote.pem)',
            default='./remote.pem',
            type=str)
    common_group.add_argument('-e', '--ciphers',
            help='Ciphers to use for encryption. '
                 'Run `openssl ciphers` to see available ciphers',
            type=str)
    common_group.add_argument('--backlog',
            help='Backlog for the listening socket (default: 128)',
            default=128,
            type=int)
    common_group.add_argument('--loglevel',
            help='Log level (default: info)',
            default='info',
            type=str,
            choices=[
                'critical', 'fatal', 'error',
                'warning', 'info', 'debug',
                ])
    common_group.add_argument('--block-size',
            help='Block size for data streaming, in bytes (default: 2048)',
            default=DEFAULT_BLOCK_SIZE,
            type=int)

    client_group = parser.add_argument_group('Client Options')

    client_group.add_argument('-p', '--peer',
            metavar='<ADDRESS>:<PORT>',
            help='Peer (server) address',
            type=str)

    server_group = parser.add_argument_group('Server Options')

    server_group.add_argument('-f', '--forward',
            metavar='<ADDRESS>:<PORT>',
            help='Simply forward all connections to the given address',
            type=str)

    return parser.parse_args()

def parse_ip_port(addr_str):
    last_col = addr_str.rfind(':')
    return (addr_str[0:last_col], int(addr_str[last_col+1:]))

def main():
    args = parse_arguments()
    logging.basicConfig(
            level=args.loglevel.upper(),
            format='%(asctime)s:%(levelname)s:%(name)s:%(message)s')

    loop = asyncio.get_event_loop()

    if args.client:
        server = client_main(loop, args)
    elif args.server:
        server = server_main(loop, args)

    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass

    server.close()
    loop.run_until_complete(server.wait_closed())
    loop.close()

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


<commit_before>"""Handles linter tasks for a given repo."""

# Copyright 2015-2016 Capstone Team G
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import logging
import os
from typing import List
from urllib.parse import urljoin
from uuid import uuid4, UUID

import github

from lintable_db.database import DatabaseHandler
from lintable_db.models import User, Repo
from lintable_git.git_handler import GitHandler
from lintable_lintball.lint_error import LintError
from lintable_lintball.lint_report import LintReport
from lintable_lintball.lint_wrapper import LintWrapper
from lintable_lintball.runner import runner
from lintable_linters.whitespace_file_linter import WhitespaceFileLinter
from lintable_processes.db_handler import DBHandler
from lintable_processes.log_handler import LogHandler
from lintable_processes.process_handler import ProcessHandler
from lintable_processes.status_handler import StatusHandler
from lintable_settings.settings import LINTWEB_SETTINGS

@runner.task(serializer='json')
def lint_github(payload: json, target_url: str, task_id: UUID):
    """Receive a task to lint a Github repo."""

    logger = logging.getLogger()
    logger.error('received payload')

    if 'action' not in payload:
        return # New registrations of repos send a payload that does not have an action.

    if payload['action'] not in ['opened', 'synchronize', 'reopened']:
        logger.error('payload ignored...')
        logger.error('payload action: {}'.format(payload['action']))
        return

    github_id = payload['repository']['owner']['id']

    owner = DatabaseHandler.get_user(github_id)

    oauth_key = owner.get_oauth_token() if isinstance(owner, User) else None

    if oauth_key is None:
        logger.error('Unable to locate oauth_token for {user} with id of {id}'.format(user=owner, id=github_id))
        return

    full_name = payload['repository']['full_name']
    repo_url = 'https://{oauth_key}@github.com/{full_name}.git'.format(
        oauth_key=oauth_key,
        full_name=full_name)

    sha1_a = payload['pull_request']['head']['sha']
    sha1_b = payload['pull_request']['base']['sha']

    repo_id = payload['repository']['id']  # type: int

    if DatabaseHandler.get_repo(repo_id) is None:
        repo = Repo(repo_id=repo_id, owner=owner,
                    url=payload['repository']['clone_url'])
        repo.save()

    target_url = target_url + '/'
    target_url = urljoin(target_url, str(task_id))

    client_id = LINTWEB_SETTINGS['github']['CLIENT_ID']
    client_secret = LINTWEB_SETTINGS['github']['CLIENT_SECRET']

    github_api = github.Github(login_or_token=oauth_key, client_id=client_id,
                               client_secret=client_secret)

    logger.error('getting repo for {full_name}'.format(full_name=full_name))

    github_repo = github_api.get_repo(full_name_or_id=full_name, lazy=False)

    logger.error('repo: {repo}'.format(repo=github_repo.id))

    github_commit = github_repo.get_commit(sha1_a)

    logger.error('target_url for status: {target_url}'.format(target_url=target_url))

    process_handler = ProcessHandler(repo=repo_url,
                                     uuid=task_id,
                                     status_handler=StatusHandler(github_commit=github_commit, target_url=target_url),
                                     logger=LogHandler(logger),
                                     db=DBHandler(repo_id=repo_id))

    git_handler = GitHandler(process_handler=process_handler,
                             repo_url=repo_url,
                             sha1_a=sha1_a,
                             sha1_b=sha1_b)

    lint_process(git_handler, process_handler)

    return

def lint_process(git_handler: GitHandler,
                 process_handler: ProcessHandler,
                 linters=None):
    """Get the files we're interested in, then start the linter."""

    if linters is None:
        linters = [WhitespaceFileLinter()]

    git_handler.started()

    git_handler.clone_repo()

    git_handler.retrieve_changed_files_from_commit()

    lintball(process_handler, linters)

    return

def lintball(handler: ProcessHandler, linters: List[LintWrapper]):
    """Run a linter or linters."""

    a_path = os.path.join(handler.local_path, 'a')
    b_path = os.path.join(handler.local_path, 'b')

    lint_errors = {}

    for filename in handler.files:
        a_file = os.path.join(a_path, filename)
        b_file = os.path.join(b_path, filename)

        a_results = lint(a_file, linters, handler) if os.path.exists(
            a_file) else []
        b_results = lint(b_file, linters, handler) if os.path.exists(
            b_file) else []
        lint_errors[filename] = [results for results in b_results if
                                 results not in a_results]

    lint_report = LintReport(errors=lint_errors)

    handler.report(lint_report)

    handler.finish()

    return

def lint(filename: str, linters: List[LintWrapper], handler: ProcessHandler) -> List[LintError]:
    """Run a linter or linters."""

    lint_errors = []

    for linter in linters:
        handler.lint_file(linter=str(linter), file=filename)
        lint_errors.extend(linter.lint(filename))

    return lint_errors
<commit_msg>fixed getting task_id from celery while not including it as a keyword arg<commit_after>"""Handles linter tasks for a given repo."""

# Copyright 2015-2016 Capstone Team G
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import logging
import os
from typing import List
from urllib.parse import urljoin
from uuid import uuid4, UUID

import celery
import github

from lintable_db.database import DatabaseHandler
from lintable_db.models import User, Repo
from lintable_git.git_handler import GitHandler
from lintable_lintball.lint_error import LintError
from lintable_lintball.lint_report import LintReport
from lintable_lintball.lint_wrapper import LintWrapper
from lintable_lintball.runner import runner
from lintable_linters.whitespace_file_linter import WhitespaceFileLinter
from lintable_processes.db_handler import DBHandler
from lintable_processes.log_handler import LogHandler
from lintable_processes.process_handler import ProcessHandler
from lintable_processes.status_handler import StatusHandler
from lintable_settings.settings import LINTWEB_SETTINGS

@runner.task(bind=True, serializer='json')
def lint_github(context, payload: json, target_url: str):
    """Receive a task to lint a Github repo."""
    task_id = context.request.id

    logger = logging.getLogger()
    logger.error('received payload')

    if 'action' not in payload:
        return  # New registrations of repos send a payload that does not have an action.

    if payload['action'] not in ['opened', 'synchronize', 'reopened']:
        logger.error('payload ignored...')
        logger.error('payload action: {}'.format(payload['action']))
        return

    github_id = payload['repository']['owner']['id']

    owner = DatabaseHandler.get_user(github_id)

    oauth_key = owner.get_oauth_token() if isinstance(owner, User) else None

    if oauth_key is None:
        logger.error('Unable to locate oauth_token for {user} with id of {id}'.format(user=owner, id=github_id))
        return

    full_name = payload['repository']['full_name']
    repo_url = 'https://{oauth_key}@github.com/{full_name}.git'.format(
        oauth_key=oauth_key,
        full_name=full_name)

    sha1_a = payload['pull_request']['head']['sha']
    sha1_b = payload['pull_request']['base']['sha']

    repo_id = payload['repository']['id']  # type: int

    if DatabaseHandler.get_repo(repo_id) is None:
        repo = Repo(repo_id=repo_id, owner=owner,
                    url=payload['repository']['clone_url'])
        repo.save()

    target_url = target_url + '/'
    target_url = urljoin(target_url, str(task_id))

    client_id = LINTWEB_SETTINGS['github']['CLIENT_ID']
    client_secret = LINTWEB_SETTINGS['github']['CLIENT_SECRET']

    github_api = github.Github(login_or_token=oauth_key, client_id=client_id,
                               client_secret=client_secret)

    logger.error('getting repo for {full_name}'.format(full_name=full_name))

    github_repo = github_api.get_repo(full_name_or_id=full_name, lazy=False)

    logger.error('repo: {repo}'.format(repo=github_repo.id))

    github_commit = github_repo.get_commit(sha1_a)

    logger.error('target_url for status: {target_url}'.format(target_url=target_url))

    process_handler = ProcessHandler(repo=repo_url,
                                     uuid=task_id,
                                     status_handler=StatusHandler(github_commit=github_commit, target_url=target_url),
                                     logger=LogHandler(logger),
                                     db=DBHandler(repo_id=repo_id))

    git_handler = GitHandler(process_handler=process_handler,
                             repo_url=repo_url,
                             sha1_a=sha1_a,
                             sha1_b=sha1_b)

    lint_process(git_handler, process_handler)

    return

def lint_process(git_handler: GitHandler,
                 process_handler: ProcessHandler,
                 linters=None):
    """Get the files we're interested in, then start the linter."""

    if linters is None:
        linters = [WhitespaceFileLinter()]

    git_handler.started()

    git_handler.clone_repo()

    git_handler.retrieve_changed_files_from_commit()

    lintball(process_handler, linters)

    return

def lintball(handler: ProcessHandler, linters: List[LintWrapper]):
    """Run a linter or linters."""

    a_path = os.path.join(handler.local_path, 'a')
    b_path = os.path.join(handler.local_path, 'b')

    lint_errors = {}

    for filename in handler.files:
        a_file = os.path.join(a_path, filename)
        b_file = os.path.join(b_path, filename)

        a_results = lint(a_file, linters, handler) if os.path.exists(
            a_file) else []
        b_results = lint(b_file, linters, handler) if os.path.exists(
            b_file) else []
        lint_errors[filename] = [results for results in b_results if
                                 results not in a_results]

    lint_report = LintReport(errors=lint_errors)

    handler.report(lint_report)

    handler.finish()

    return

def lint(filename: str, linters: List[LintWrapper], handler: ProcessHandler) -> List[LintError]:
    """Run a linter or linters."""

    lint_errors = []

    for linter in linters:
        handler.lint_file(linter=str(linter), file=filename)
        lint_errors.extend(linter.lint(filename))

    return lint_errors
<|endoftext|>


<commit_before>#!/usr/bin/env python

# Dependency Generator for C/C++ and Fortran
#
# Language: Python (2.7 or 3.x)
#
# Version: 1.0.0
#
# Copyright (c) 2000-2018 Objexx Engineering, Inc. All Rights Reserved.

# Notes:
# . Dependency files for use with GNU make
# . Only processes files with supported source C/C++ extensions
# . Leaves specified paths on prerequisites but does not add path where found
# . Adds a line with all prerequisites as targets to avoid make error if one is removed or renamed
# . Include search path defaults to environment variables:
#   . C/C++: CPATH then CPLUS_INCLUDE_PATH then INCLUDE
#   . Fortran: CPATH then CPLUS_INCLUDE_PATH then INCLUDE then FINCLUDE
# . Use --no-mod for compilers like GFortran that don't regenerate unchanged .mod files

# Imports
import argparse, os, re

# Globals
C_ext = re.compile( '(c|cc|cpp|cxx|c\+\+|h|hh|hpp|hxx|h\+\+|ii|ipp|ixx|i\+\+)', re.I )
C_inc = re.compile( '\s*#\s*include +<([^<>]+\.[^<>]+)>' ) # Exclude <name> C++ Standard Library headers
C_inq = re.compile( '\s*#\s*include +"([^"]+)"' )
C_sys = re.compile( 'sys/' ) # Exclude <sys/header> C/C++ system headers
C_path = None
F_ext = re.compile( '(f|for|f90|f9x|f95|f03|f08|fh|fi|inc)', re.I )
F_inc = re.compile( "\s*INCLUDE +'([^']+)'", re.I )
F_mod = re.compile( "\s*MODULE +([^ \t\n]+)", re.I )
F_use = re.compile( "\s*USE +([^ ,\t\n]+)", re.I )
F_sys = [ 'ISO_C_BINDING', 'ISO_FORTRAN_ENV', 'IEEE_EXCEPTIONS', 'IEEE_ARITHMETIC', 'IEEE_FEATURES', 'OMP_LIB', 'OMP_LIB_KINDS', 'OPENACC', 'HDF5', 'H5GLOBAL' ]
F_path = None

# Main
def main():

    # Get options and arguments
    parser = argparse.ArgumentParser( description = 'Generates C/C++ and Fortran dependency files' )
    parser.add_argument( '--inc', help = 'include path or env var [CPATH|CPLUS_INCLUDE_PATH|INCLUDE|FINCLUDE]' )
    parser.add_argument( '--ext', help = 'object file extension [obj|o]', default = ( 'obj' if os.name == 'nt' else 'o' ) )
    parser.add_argument( '--no-mod', dest = 'no_mod', help = 'Suppress mod file targets', action = 'store_true', default = False )
    parser.add_argument( 'source', help = 'Source file' )
    arg = parser.parse_args()

    # Set up include search paths
    global C_path
    if not arg.inc:
        if 'CPATH' in os.environ:
            C_path = os.environ[ 'CPATH' ]
        elif 'CPLUS_INCLUDE_PATH' in os.environ:
            C_path = os.environ[ 'CPLUS_INCLUDE_PATH' ]
        elif 'INCLUDE' in os.environ:
            C_path = os.environ[ 'INCLUDE' ]
    elif arg.inc in os.environ: # Treat arg as env var name
        C_path = os.environ[ arg.inc ]
    else: # Treat arg as path
        C_path = arg.inc
    C_path = C_path.split( os.pathsep ) if C_path else None
    global F_path
    if not arg.inc:
        if 'CPATH' in os.environ:
            F_path = os.environ[ 'CPATH' ]
        elif 'CPLUS_INCLUDE_PATH' in os.environ:
            F_path = os.environ[ 'CPLUS_INCLUDE_PATH' ]
        elif 'INCLUDE' in os.environ:
            F_path = os.environ[ 'INCLUDE' ]
        elif 'FINCLUDE' in os.environ:
            F_path = os.environ[ 'FINCLUDE' ]
    elif arg.inc in os.environ: # Treat arg as env var name
        F_path = os.environ[ arg.inc ]
    else: # Treat arg as path
        F_path = arg.inc
    F_path = F_path.split( os.pathsep ) if F_path else None

    # Parse source file name
    if not os.path.isfile( arg.source ):
        raise IOError( 'Source file not found: ' + str( arg.source ) )
    src_name = os.path.basename( arg.source )
    src_base, src_ext = os.path.splitext( src_name )
    dep_name = src_base + '.d'
    obj_name = src_base + '.' + arg.ext
    if src_ext: src_ext = src_ext[ 1: ]
    if not ( src_base and src_ext and ( C_ext.match( src_ext ) or F_ext.match( src_ext ) ) ):
        raise ValueError( 'Not a recognized C/C++ or Fortran file name extension: ' + str( src_ext ) )

    # Build dependencies
    tar = obj_name + ' ' + dep_name
    dep = [ src_name ]
    if C_ext.match( src_ext ):
        dep_list = list( C_deps( os.path.abspath( arg.source ), add = False ) )
        mod_list = None
        use_list = None
    elif F_ext.match( src_ext ):
        dep_list, mod_list, use_list = F_deps( os.path.abspath( arg.source ), add = False )
        dep_list = list( dep_list )
        mod_list = list( mod_list )
        mod_list.sort()
        use_list = list( use_list )
        use_list.sort()
    else:
        assert False, 'Not a recognized C/C++ or Fortran file name extension: ' + str( src_ext )
    dep_list.sort()
    dep.extend( dep_list )
    dep_str = ' '.join( dep )
    if mod_list and ( not arg.no_mod ):
        for mod in mod_list:
            tar += ' ' + mod + '.mod' # Assumes lowercase module naming
    if use_list:
        for use in use_list:
            dep_str += ' ' + use + '.mod' # Assumes lowercase module naming

    # Write dependency file
    dep_file = open( dep_name, 'w' )
    dep_file.write( tar + ' : ' + dep_str + '\n' )
    dep_file.write( dep_str + ' :\n' )
    if mod_list: # Add rules for building .mod files: Implicit rules don't work because .mod file names != source file names
        if arg.no_mod:
            for mod in mod_list:
                dep_file.write( mod + '.mod : ' + obj_name + '\n' )
                dep_file.write( '\t@true\n' )
        else:
            for mod in mod_list:
                dep_file.write( mod + '.mod : ' + src_name + '\n' )
                if os.name == 'nt':
                    dep_file.write( '\t$(FC) $(FFLAGS) -c $(subst /,\,$<)\n' )
                else:
                    dep_file.write( '\t$(FC) $(FFLAGS) -c $<\n' )
    dep_file.close()

def C_deps( fname, fdeps = None, par_dir = None, quoted = False, add = True ):
    '''C/C++ dependencies of a file: Recursive'''
    if fdeps is None: fdeps = set()
    if par_dir is None: par_dir = []
    try:
        gname = None
        if ( not os.path.isabs( fname ) ) and ( not os.path.isfile( fname ) ): # Search include path list
            if quoted: # Check relative to parent directories
                for adir in par_dir:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
            if not gname: # Use include search path
                for adir in C_path:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
        if gname is None: gname = fname
        dfile = open( gname, 'rU' )
        if add: fdeps.add( fname ) # Only add dependency if it is found
        par_new = os.path.dirname( os.path.abspath( gname ) )
        if ( not par_dir ) or ( par_dir[ 0 ] != par_new ): # Push parent dir onto front of list
            par_dir.insert( 0, par_new )
        else: # Flag not to pop since didn't push
            par_new = None
        for line in dfile:
            m = C_inc.match( line ) # #include <header> form
            if m :
                dep = m.group( 1 )
                if dep and ( dep not in fdeps ) and ( not C_sys.match( dep ) ):
                    C_deps( dep, fdeps, par_dir ) # Recurse
            else:
                m = C_inq.match( line ) # #include "header" form
                if m:
                    dep = m.group( 1 )
                    if dep and ( dep not in fdeps ) and ( not C_sys.match( dep ) ):
                        C_deps( dep, fdeps, par_dir, quoted = True ) # Recurse
        dfile.close()
        if par_new: del par_dir[ 0 ] # Pop parent dir from front of list
    except:
        pass # Skip missing file
    return fdeps

def F_deps( fname, fdeps = None, fmods = None, fuses = None, par_dir = None, quoted = False, add = True ):
    '''Fortran dependencies of a file: Recursive'''
    if fdeps is None: fdeps = set()
    if fmods is None: fmods = set()
    if fuses is None: fuses = set()
    if par_dir is None: par_dir = []
    try:
        gname = None
        if ( not os.path.isabs( fname ) ) and ( not os.path.isfile( fname ) ): # Search include path list
            if quoted: # Check relative to parent directories
                for adir in par_dir:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
            if not gname: # Use include search path
                for adir in F_path:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
        if gname is None: gname = fname
        dfile = open( gname, 'rU' )
        if add: fdeps.add( fname ) # Only add dependency if it is found
        par_new = os.path.dirname( os.path.abspath( gname ) )
        if ( not par_dir ) or ( par_dir[ 0 ] != par_new ): # Push parent dir onto front of list
            par_dir.insert( 0, par_new )
        else: # Flag not to pop since didn't push
            par_new = None
        for line in dfile:
            m = F_inc.match( line ) # INCLUDE 'header'
            if m:
                dep = m.group( 1 )
                if dep and ( dep not in fdeps ):
                    F_deps( dep, fdeps, fmods, fuses, par_dir ) # Recurse
            else:
                m = F_mod.match( line ) # MODULE modname
                if m:
                    mod = m.group( 1 ).lower()
                    if mod and ( mod not in fmods ) and ( mod != 'procedure' ):
                        fmods.add( mod )
                else:
                    m = F_use.match( line ) # USE modname
                    if m:
                        mod = m.group( 1 ).lower()
                        if mod.upper() not in F_sys:
                            if mod and ( mod not in fuses ):
                                fuses.add( mod )
        dfile.close()
        if par_new: del par_dir[ 0 ] # Pop parent dir from front of list
    except:
        pass # Skip missing file
    return fdeps, fmods, fuses

# Runner
if __name__ == '__main__':
    main()
<commit_msg>makedep.py Py3 update<commit_after>#!/usr/bin/env python

# Dependency Generator for C/C++ and Fortran
#
# Language: Python (2.7 or 3.x)
#
# Version: 1.0.0
#
# Copyright (c) 2000-2018 Objexx Engineering, Inc. All Rights Reserved.

# Notes:
# . Dependency files for use with GNU make
# . Only processes files with supported source C/C++ extensions
# . Leaves specified paths on prerequisites but does not add path where found
# . Adds a line with all prerequisites as targets to avoid make error if one is removed or renamed
# . Include search path defaults to environment variables:
#   . C/C++: CPATH then CPLUS_INCLUDE_PATH then INCLUDE
#   . Fortran: CPATH then CPLUS_INCLUDE_PATH then INCLUDE then FINCLUDE
# . Use --no-mod for compilers like GFortran that don't regenerate unchanged .mod files

# Imports
import argparse, os, re, sys

# Globals
C_ext = re.compile( '(c|cc|cpp|cxx|c\+\+|h|hh|hpp|hxx|h\+\+|ii|ipp|ixx|i\+\+)', re.I )
C_inc = re.compile( '\s*#\s*include +<([^<>]+\.[^<>]+)>' ) # Exclude <name> C++ Standard Library headers
C_inq = re.compile( '\s*#\s*include +"([^"]+)"' )
C_sys = re.compile( 'sys/' ) # Exclude <sys/header> C/C++ system headers
C_path = None
F_ext = re.compile( '(f|for|f90|f9x|f95|f03|f08|fh|fi|inc)', re.I )
F_inc = re.compile( "\s*INCLUDE +'([^']+)'", re.I )
F_mod = re.compile( "\s*MODULE +([^ \t\n]+)", re.I )
F_use = re.compile( "\s*USE +([^ ,\t\n]+)", re.I )
F_sys = [ 'ISO_C_BINDING', 'ISO_FORTRAN_ENV', 'IEEE_EXCEPTIONS', 'IEEE_ARITHMETIC', 'IEEE_FEATURES', 'OMP_LIB', 'OMP_LIB_KINDS', 'OPENACC', 'HDF5', 'H5GLOBAL' ]
F_path = None

# Main
def main():

    # Get options and arguments
    parser = argparse.ArgumentParser( description = 'Generates C/C++ and Fortran dependency files' )
    parser.add_argument( '--inc', help = 'include path or env var [CPATH|CPLUS_INCLUDE_PATH|INCLUDE|FINCLUDE]' )
    parser.add_argument( '--ext', help = 'object file extension [obj|o]', default = ( 'obj' if os.name == 'nt' else 'o' ) )
    parser.add_argument( '--no-mod', dest = 'no_mod', help = 'Suppress mod file targets', action = 'store_true', default = False )
    parser.add_argument( 'source', help = 'Source file' )
    arg = parser.parse_args()

    # Set up include search paths
    global C_path
    if not arg.inc:
        if 'CPATH' in os.environ:
            C_path = os.environ[ 'CPATH' ]
        elif 'CPLUS_INCLUDE_PATH' in os.environ:
            C_path = os.environ[ 'CPLUS_INCLUDE_PATH' ]
        elif 'INCLUDE' in os.environ:
            C_path = os.environ[ 'INCLUDE' ]
    elif arg.inc in os.environ: # Treat arg as env var name
        C_path = os.environ[ arg.inc ]
    else: # Treat arg as path
        C_path = arg.inc
    C_path = C_path.split( os.pathsep ) if C_path else None
    global F_path
    if not arg.inc:
        if 'CPATH' in os.environ:
            F_path = os.environ[ 'CPATH' ]
        elif 'CPLUS_INCLUDE_PATH' in os.environ:
            F_path = os.environ[ 'CPLUS_INCLUDE_PATH' ]
        elif 'INCLUDE' in os.environ:
            F_path = os.environ[ 'INCLUDE' ]
        elif 'FINCLUDE' in os.environ:
            F_path = os.environ[ 'FINCLUDE' ]
    elif arg.inc in os.environ: # Treat arg as env var name
        F_path = os.environ[ arg.inc ]
    else: # Treat arg as path
        F_path = arg.inc
    F_path = F_path.split( os.pathsep ) if F_path else None

    # Parse source file name
    if not os.path.isfile( arg.source ):
        raise IOError( 'Source file not found: ' + str( arg.source ) )
    src_name = os.path.basename( arg.source )
    src_base, src_ext = os.path.splitext( src_name )
    dep_name = src_base + '.d'
    obj_name = src_base + '.' + arg.ext
    if src_ext: src_ext = src_ext[ 1: ]
    if not ( src_base and src_ext and ( C_ext.match( src_ext ) or F_ext.match( src_ext ) ) ):
        raise ValueError( 'Not a recognized C/C++ or Fortran file name extension: ' + str( src_ext ) )

    # Build dependencies
    tar = obj_name + ' ' + dep_name
    dep = [ src_name ]
    if C_ext.match( src_ext ):
        dep_list = list( C_deps( os.path.abspath( arg.source ), add = False ) )
        mod_list = None
        use_list = None
    elif F_ext.match( src_ext ):
        dep_list, mod_list, use_list = F_deps( os.path.abspath( arg.source ), add = False )
        dep_list = list( dep_list )
        mod_list = list( mod_list )
        mod_list.sort()
        use_list = list( use_list )
        use_list.sort()
    else:
        assert False, 'Not a recognized C/C++ or Fortran file name extension: ' + str( src_ext )
    dep_list.sort()
    dep.extend( dep_list )
    dep_str = ' '.join( dep )
    if mod_list and ( not arg.no_mod ):
        for mod in mod_list:
            tar += ' ' + mod + '.mod' # Assumes lowercase module naming
    if use_list:
        for use in use_list:
            dep_str += ' ' + use + '.mod' # Assumes lowercase module naming

    # Write dependency file
    dep_file = open( dep_name, 'w' )
    dep_file.write( tar + ' : ' + dep_str + '\n' )
    dep_file.write( dep_str + ' :\n' )
    if mod_list: # Add rules for building .mod files: Implicit rules don't work because .mod file names != source file names
        if arg.no_mod:
            for mod in mod_list:
                dep_file.write( mod + '.mod : ' + obj_name + '\n' )
                dep_file.write( '\t@true\n' )
        else:
            for mod in mod_list:
                dep_file.write( mod + '.mod : ' + src_name + '\n' )
                if os.name == 'nt':
                    dep_file.write( '\t$(FC) $(FFLAGS) -c $(subst /,\,$<)\n' )
                else:
                    dep_file.write( '\t$(FC) $(FFLAGS) -c $<\n' )
    dep_file.close()

def C_deps( fname, fdeps = None, par_dir = None, quoted = False, add = True ):
    '''C/C++ dependencies of a file: Recursive'''
    if fdeps is None: fdeps = set()
    if par_dir is None: par_dir = []
    try:
        gname = None
        if ( not os.path.isabs( fname ) ) and ( not os.path.isfile( fname ) ): # Search include path list
            if quoted: # Check relative to parent directories
                for adir in par_dir:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
            if not gname: # Use include search path
                for adir in C_path:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
        if gname is None: gname = fname
        if sys.version_info >= ( 3, 0 ):
            dfile = open( gname, 'r', newline = None )
        else:
            dfile = open( gname, 'rU' )
        if add: fdeps.add( fname ) # Only add dependency if it is found
        par_new = os.path.dirname( os.path.abspath( gname ) )
        if ( not par_dir ) or ( par_dir[ 0 ] != par_new ): # Push parent dir onto front of list
            par_dir.insert( 0, par_new )
        else: # Flag not to pop since didn't push
            par_new = None
        for line in dfile:
            m = C_inc.match( line ) # #include <header> form
            if m :
                dep = m.group( 1 )
                if dep and ( dep not in fdeps ) and ( not C_sys.match( dep ) ):
                    C_deps( dep, fdeps, par_dir ) # Recurse
            else:
                m = C_inq.match( line ) # #include "header" form
                if m:
                    dep = m.group( 1 )
                    if dep and ( dep not in fdeps ) and ( not C_sys.match( dep ) ):
                        C_deps( dep, fdeps, par_dir, quoted = True ) # Recurse
        dfile.close()
        if par_new: del par_dir[ 0 ] # Pop parent dir from front of list
    except:
        pass # Skip missing file
    return fdeps

def F_deps( fname, fdeps = None, fmods = None, fuses = None, par_dir = None, quoted = False, add = True ):
    '''Fortran dependencies of a file: Recursive'''
    if fdeps is None: fdeps = set()
    if fmods is None: fmods = set()
    if fuses is None: fuses = set()
    if par_dir is None: par_dir = []
    try:
        gname = None
        if ( not os.path.isabs( fname ) ) and ( not os.path.isfile( fname ) ): # Search include path list
            if quoted: # Check relative to parent directories
                for adir in par_dir:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
            if not gname: # Use include search path
                for adir in F_path:
                    tname = os.path.join( adir, fname )
                    if os.path.isfile( tname ):
                        if not ' ' in tname: gname = tname # Skip dependencies in paths with spaces for GNU Make compatibility
                        break
        if gname is None: gname = fname
        if sys.version_info >= ( 3, 0 ):
            dfile = open( gname, 'r', newline = None )
        else:
            dfile = open( gname, 'rU' )
        if add: fdeps.add( fname ) # Only add dependency if it is found
        par_new = os.path.dirname( os.path.abspath( gname ) )
        if ( not par_dir ) or ( par_dir[ 0 ] != par_new ): # Push parent dir onto front of list
            par_dir.insert( 0, par_new )
        else: # Flag not to pop since didn't push
            par_new = None
        for line in dfile:
            m = F_inc.match( line ) # INCLUDE 'header'
            if m:
                dep = m.group( 1 )
                if dep and ( dep not in fdeps ):
                    F_deps( dep, fdeps, fmods, fuses, par_dir ) # Recurse
            else:
                m = F_mod.match( line ) # MODULE modname
                if m:
                    mod = m.group( 1 ).lower()
                    if mod and ( mod not in fmods ) and ( mod != 'procedure' ):
                        fmods.add( mod )
                else:
                    m = F_use.match( line ) # USE modname
                    if m:
                        mod = m.group( 1 ).lower()
                        if mod.upper() not in F_sys:
                            if mod and ( mod not in fuses ):
                                fuses.add( mod )
        dfile.close()
        if par_new: del par_dir[ 0 ] # Pop parent dir from front of list
    except:
        pass # Skip missing file
    return fdeps, fmods, fuses

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