<commit_before>from abc import ABCMeta, abstractmethod
import Monstr.Core.DB as DB
import Monstr.Core.Utils as Utils
import Monstr.Core.Constants as Constants

# ,----------------------.
# |BaseNodule            |
# |----------------------|
# |+string name          |
# |+obj table_schemas    |
# |----------------------|
# |+void Initialize()    |
# |+obj PrepareRetrieve()|
# |+obj Retrieve()       |
# |+obj InsertToDB()     |
# |+obj Analyze()        |
# |+obj React()          |
# |+obj Run()            |
# `----------------------'

class BaseModule():
    __metaclass__ = ABCMeta

    name = None
    table_schemas = None
    tables = None
    status_table = None
    status_list = []
    events_table = None

    rest_links = {}

    db_handler = None

    status_schema = {'status': (DB.Column('id', DB.Integer, primary_key=True),
                                DB.Column('name', DB.String(64)),
                                DB.Column('status', DB.Integer),
                                DB.Column('time', DB.DateTime(True)),
                                DB.Column('description', DB.Text),)}

    journal_schema = (DB.Column('id', DB.Integer, primary_key=True),
                      DB.Column('module', DB.String(64)),
                      DB.Column('time', DB.DateTime(True)),
                      DB.Column('result', DB.String(32)),
                      DB.Column('step', DB.String(32)),
                      DB.Column('description', DB.Text),)

    events_schema = (DB.Column('id', DB.BigInteger, primary_key=True),
                     DB.Column('module', DB.String(64)),
                     DB.Column('name', DB.String(64)),
                     DB.Column('type', DB.String(32)),
                     DB.Column('time', DB.DateTime(True)),
                     DB.Column('severity', DB.Integer),
                     DB.Column('description', DB.Text),)

    # ==========================================================================
    # Database functions
    # ==========================================================================

    def _db_incert_journal_row(self, row):
        self.db_handler.insert(self.journal, row)

    def _db_incert_event_row(self, row):
        self.db_handler.insert(self.events_table, row)

    def _db_get_status_table_repr(self):
        return [x._asdict() for x in self.db_handler.get_session().query(self.status_table['status']).all()]

    def _db_update_status(self, statuses):
        conn = self.db_handler.get_engine().connect()
        for status in statuses:
            update = self.status_table['status'].update().values(status=status['status'], time=status['time'], description=status['description']).where(self.status_table['status'].c.name == status['name'])
            conn.execute(update)

    # ==========================================================================
    # Common functions
    # ==========================================================================

    def _create_journal_row(self, result, step=None, description=None):
        row = {'module': self.name,
               'time': Utils.get_UTC_now(),
               'result': result,
               'step': step,
               'description': description}
        return row

    def write_to_journal(self, result, step=None, description=None):
        row = self._create_journal_row(result, step, description)
        self._db_incert_journal_row(row)

    def write_error_to_journal(self, result, step=None, error=None):
        description = (type(error).__name__ + ': ' + error.message) if error is not None else None
        self.write_to_journal(result, step, description)

    # --------------------------------------------------------------------------

    def create_event(self, name, event_type, time, severity, description):
        event = {'module': self.name,
                 'name': name,
                 'type': event_type,
                 'time': time,
                 'severity': severity,
                 'description': description}
        return event

    def write_event(self, event):
        self._db_incert_event_row(event)

    # --------------------------------------------------------------------------

    def _create_params(self, default_params, params):
        result = {}
        for key in default_params:
            if key not in params:
                result[key] = default_params[key]
            else:
                result[key] = type(default_params[key])(params[key])
        return result

    def get_last_status(self):
        return self._db_get_status_table_repr()

    def update_status(self, current_statuses):
        previous_status = self._db_get_status_table_repr()
        last_status = dict([(str(x['name']), int(x['status'])) for x in previous_status])
        update_list = []
        event_list = []
        for status in current_statuses:
            if last_status[status['name']] != status['status']:
                update_list.append(status)
                # Generate event and write to DB
                event_name = status['name'] + ':' + Constants.STATUS[last_status[status['name']]] + '->' + Constants.STATUS[status['status']]
                event = self.create_event(event_name, 'StatusChange', status['time'], status['status'], status['description'])
                event_list.append(event)
                self.write_event(event)
        self._db_update_status(update_list)
        return event_list

    # --------------------------------------------------------------------------

    def __init__(self):
        self.db_handler = DB.DBHandler()
        self.rest_links = {'getModuleStatus': self.GetModuleStatus}

    def Initialize(self):
        if self.name is None:
            raise "Module require name"
        if self.table_schemas is None:
            raise "Module require schemas list"
        self.tables = self.db_handler.initialize(self.table_schemas, self.name)
        self.status_table = self.db_handler.initialize(self.status_schema, self.name, self.status_list)
        self.journal = self.db_handler.getOrCreateTable('monstr_Journal', self.journal_schema)
        self.events_table = self.db_handler.getOrCreateTable('monstr_Events', self.events_schema)

    def PrepareRetrieve(self):
        return {}

    def Retrieve(self, params):
        pass

    def Analyze(self, data):
        return []

    def React(self, events):
        pass

    def InsertToDB(self, data):
        for schema in data:
            table = self.tables[schema]
            self.db_handler.bulk_insert(table, data[schema])
        return

    def ExecuteCheck(self):
        try:
            self.Initialize()
        except Exception as e:
            self.write_error_to_journal('Fail', 'Initialize', e)
            print e
            return

        try:
            params = self.PrepareRetrieve()
        except Exception as e:
            self.write_error_to_journal('Fail', 'PrepareRetrieve', e)
            print e
            return

        try:
            data = self.Retrieve(params)
        except Exception as e:
            self.write_error_to_journal('Fail', 'Retrieve', e)
            print e
            return

        try:
            self.InsertToDB(data)
        except Exception as e:
            self.write_error_to_journal('Fail', 'InsertToDB', e)
            print e
            return

        #try:
        events = self.Analyze(data)
        #except Exception as e:
        #    self.write_error_to_journal('Fail', 'Analyze', e)
        #    print 'Analyze error'
        #    print e
        #    return

        try:
            self.React(data)
        except Exception as e:
            self.write_error_to_journal('Fail', 'React', e)
            print e
            return

        self.write_to_journal('Success')

    # ==========================================================================
    #                 Web
    # ==========================================================================

    def GetModuleStatus(self, incoming_params):
        response = {}
        params = incoming_params
        try:
            result = self._db_get_status_table_repr()
            response = {'data': result,
                        'applied_params': params,
                        'success': True}
        except Exception as e:
            response = {'data': {},
                        'incoming_params': incoming_params,
                        'success': False,
                        'error': type(e).__name__ + ': ' + e.message,
                        'description': 'Error inside BaseModule.GetModuleStatus'}

        return response
<commit_msg>FIX: old codelines removed<commit_after>from abc import ABCMeta, abstractmethod
import Monstr.Core.DB as DB
import Monstr.Core.Utils as Utils
import Monstr.Core.Constants as Constants

# ,----------------------.
# |BaseNodule            |
# |----------------------|
# |+string name          |
# |+obj table_schemas    |
# |----------------------|
# |+void Initialize()    |
# |+obj PrepareRetrieve()|
# |+obj Retrieve()       |
# |+obj InsertToDB()     |
# |+obj Analyze()        |
# |+obj React()          |
# |+obj Run()            |
# `----------------------'

class BaseModule():
    __metaclass__ = ABCMeta

    name = None
    table_schemas = None
    tables = None
    status_table = None
    status_list = []
    journal = None
    events_table = None

    rest_links = {}

    db_handler = None

    status_schema = {'status': (DB.Column('id', DB.Integer, primary_key=True),
                                DB.Column('name', DB.String(64)),
                                DB.Column('status', DB.Integer),
                                DB.Column('time', DB.DateTime(True)),
                                DB.Column('description', DB.Text),)}

    journal_schema = (DB.Column('id', DB.Integer, primary_key=True),
                      DB.Column('module', DB.String(64)),
                      DB.Column('time', DB.DateTime(True)),
                      DB.Column('result', DB.String(32)),
                      DB.Column('step', DB.String(32)),
                      DB.Column('description', DB.Text),)

    events_schema = (DB.Column('id', DB.BigInteger, primary_key=True),
                     DB.Column('module', DB.String(64)),
                     DB.Column('name', DB.String(64)),
                     DB.Column('type', DB.String(32)),
                     DB.Column('time', DB.DateTime(True)),
                     DB.Column('severity', DB.Integer),
                     DB.Column('description', DB.Text),)

    # ==========================================================================
    # Database functions
    # ==========================================================================

    def _db_incert_journal_row(self, row):
        self.db_handler.insert(self.journal, row)

    def _db_incert_event_row(self, row):
        self.db_handler.insert(self.events_table, row)

    def _db_get_status_table_repr(self):
        return [x._asdict() for x in self.db_handler.get_session().query(self.status_table['status']).all()]

    def _db_update_status(self, statuses):
        conn = self.db_handler.get_engine().connect()
        for status in statuses:
            update = self.status_table['status'].update().values(status=status['status'], time=status['time'], description=status['description']).where(self.status_table['status'].c.name == status['name'])
            conn.execute(update)

    # ==========================================================================
    # Common functions
    # ==========================================================================

    def _create_journal_row(self, result, step=None, description=None):
        row = {'module': self.name,
               'time': Utils.get_UTC_now(),
               'result': result,
               'step': step,
               'description': description}
        return row

    def write_to_journal(self, result, step=None, description=None):
        row = self._create_journal_row(result, step, description)
        self._db_incert_journal_row(row)

    def write_error_to_journal(self, result, step=None, error=None):
        description = (type(error).__name__ + ': ' + error.message) if error is not None else None
        self.write_to_journal(result, step, description)

    # --------------------------------------------------------------------------

    def create_event(self, name, event_type, time, severity, description):
        event = {'module': self.name,
                 'name': name,
                 'type': event_type,
                 'time': time,
                 'severity': severity,
                 'description': description}
        return event

    def write_event(self, event):
        self._db_incert_event_row(event)

    # --------------------------------------------------------------------------

    def _create_params(self, default_params, params):
        result = {}
        for key in default_params:
            if key not in params:
                result[key] = default_params[key]
            else:
                result[key] = type(default_params[key])(params[key])
        return result

    def get_last_status(self):
        return self._db_get_status_table_repr()

    def update_status(self, current_statuses):
        previous_status = self._db_get_status_table_repr()
        last_status = dict([(str(x['name']), int(x['status'])) for x in previous_status])
        update_list = []
        event_list = []
        for status in current_statuses:
            if last_status[status['name']] != status['status']:
                update_list.append(status)
                # Generate event and write to DB
                event_name = status['name'] + ':' + Constants.STATUS[last_status[status['name']]] + '->' + Constants.STATUS[status['status']]
                event = self.create_event(event_name, 'StatusChange', status['time'], status['status'], status['description'])
                event_list.append(event)
                self.write_event(event)
        self._db_update_status(update_list)
        return event_list

    # --------------------------------------------------------------------------

    def __init__(self):
        self.db_handler = DB.DBHandler()
        self.rest_links = {'getModuleStatus': self.GetModuleStatus}
        self.journal = self.db_handler.getOrCreateTable('monstr_Journal', self.journal_schema)
        self.events_table = self.db_handler.getOrCreateTable('monstr_Events', self.events_schema)

    def Initialize(self):
        if self.name is None:
            raise "Module require name"
        if self.table_schemas is None:
            raise "Module require schemas list"
        self.tables = self.db_handler.initialize(self.table_schemas, self.name)
        self.status_table = self.db_handler.initialize(self.status_schema, self.name, self.status_list)

    def PrepareRetrieve(self):
        return {}

    def Retrieve(self, params):
        pass

    def Analyze(self, data):
        return []

    def React(self, events):
        pass

    def InsertToDB(self, data):
        for schema in data:
            table = self.tables[schema]
            self.db_handler.bulk_insert(table, data[schema])
        return

    def ExecuteCheck(self):
        try:
            self.Initialize()
        except Exception as e:
            self.write_error_to_journal('Fail', 'Initialize', e)
            print e
            return

        try:
            params = self.PrepareRetrieve()
        except Exception as e:
            self.write_error_to_journal('Fail', 'PrepareRetrieve', e)
            print e
            return

        try:
            data = self.Retrieve(params)
        except Exception as e:
            self.write_error_to_journal('Fail', 'Retrieve', e)
            print e
            return

        try:
            self.InsertToDB(data)
        except Exception as e:
            self.write_error_to_journal('Fail', 'InsertToDB', e)
            print e
            return

        #try:
        events = self.Analyze(data)
        #except Exception as e:
        #    self.write_error_to_journal('Fail', 'Analyze', e)
        #    print 'Analyze error'
        #    print e
        #    return

        try:
            self.React(data)
        except Exception as e:
            self.write_error_to_journal('Fail', 'React', e)
            print e
            return

        self.write_to_journal('Success')

    # ==========================================================================
    #                 Web
    # ==========================================================================

    def GetModuleStatus(self, incoming_params):
        response = {}
        params = incoming_params
        try:
            result = self._db_get_status_table_repr()
            response = {'data': result,
                        'applied_params': params,
                        'success': True}
        except Exception as e:
            response = {'data': {},
                        'incoming_params': incoming_params,
                        'success': False,
                        'error': type(e).__name__ + ': ' + e.message,
                        'description': 'Error inside BaseModule.GetModuleStatus'}

        return response
<|endoftext|>


<commit_before># Copyright (c) 2016-2017, Matteo Cafasso
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.

# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS
# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
# OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
# OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Analyse disk content to extract File System event timelines."""

import ntpath
import logging
from datetime import timedelta
from functools import lru_cache
from itertools import chain, groupby
from tempfile import NamedTemporaryFile
from collections import defaultdict, namedtuple

from vminspect.filesystem import FileSystem
from vminspect.usnjrnl import CorruptedUsnRecord, usn_journal

class FSTimeline:
    def __init__(self, disk):
        self._disk = disk
        self._filesystem = None
        self._filetype_cache = {}
        self._checksum_cache = {}
        self.logger = logging.getLogger(
            "%s.%s" % (self.__module__, self.__class__.__name__))

    def __enter__(self):
        self._filesystem = FileSystem(self._disk)
        self._filesystem.mount()

        return self

    def __exit__(self, *_):
        self._filesystem.umount()

    def __getattr__(self, attr):
        return getattr(self._filesystem, attr)

    def timeline(self):
        self.logger.debug("Extracting File System timeline events.")
        events = tuple(Event(d.inode, d.path, d.size, d.allocated, t, r)
                       for d in self._visit_filesystem()
                       for t, r in ((d.atime, 'access'),
                                    (d.mtime, 'change'),
                                    (d.ctime, 'attribute_change'),
                                    (d.crtime, 'creation'))
                       if t > 0)

        self.logger.debug("Sorting File System timeline events.")
        return sorted(events, key=lambda e: e.timestamp)

    @lru_cache(maxsize=None)
    def file(self, path):
        """Identifies the file type.

        Caches the result to reduce overhead on duplicated events.

        """
        return self._filesystem.file(path)

    @lru_cache(maxsize=None)
    def checksum(self, path):
        """Identifies the file type.

        Caches the result to reduce overhead on duplicated events.

        """
        return self._filesystem.checksum(path)

    def _visit_filesystem(self):
        """Walks through the filesystem content."""
        self.logger.debug("Parsing File System content.")

        root_partition = self._filesystem.inspect_get_roots()[0]

        yield from self._root_dirent()

        for entry in self._filesystem.filesystem_walk(root_partition):
            yield Dirent(
                entry['tsk_inode'],
                self._filesystem.path('/' + entry['tsk_name']),
                entry['tsk_size'], entry['tsk_type'],
                True if entry['tsk_flags'] & TSK_ALLOC else False,
                timestamp(entry['tsk_atime_sec'], entry['tsk_atime_nsec']),
                timestamp(entry['tsk_mtime_sec'], entry['tsk_mtime_nsec']),
                timestamp(entry['tsk_ctime_sec'], entry['tsk_ctime_nsec']),
                timestamp(entry['tsk_crtime_sec'], entry['tsk_crtime_nsec']))

    def _root_dirent(self):
        """Returns the root folder dirent as filesystem_walk API doesn't."""
        fstat = self._filesystem.stat('/')

        yield Dirent(fstat['ino'], self._filesystem.path('/'),
                     fstat['size'], 'd', True,
                     timestamp(fstat['atime'], 0),
                     timestamp(fstat['mtime'], 0),
                     timestamp(fstat['ctime'], 0),
                     0)

class NTFSTimeline(FSTimeline):
    """Inspect NTFS filesystem in order to extract a timeline of events
    containing the information related to files/directories changes.

    This feature depends on a special build of Libguestfs available at:
      https://github.com/noxdafox/libguestfs/tree/forensics

    """
    def __init__(self, disk):
        super().__init__(disk)

    def __enter__(self):
        super().__enter__()

        if self._filesystem.osname != 'windows':
            self._filesystem.umount()

            raise RuntimeError("Filesystem is not NTFS.")

        return self

    def usnjrnl_timeline(self):
        """Iterates over the changes occurred within the filesystem.

        Yields UsnJrnlEvent namedtuples containing:

            file_reference_number: known in Unix FS as inode.
            path: full path of the file.
            size: size of the file in bytes if recoverable.
            allocated: whether the file exists or it has been deleted.
            timestamp: timespamp of the change.
            changes: list of changes applied to the file.
            attributes: list of file attributes.

        """
        content = defaultdict(list)

        self.logger.debug("Extracting Update Sequence Number journal.")
        journal = self._read_journal()

        for dirent in self._visit_filesystem():
            content[dirent.inode].append(dirent)

        self.logger.debug("Generating timeline.")
        yield from generate_timeline(journal, content)

    def _read_journal(self):
        """Extracts the USN journal from the disk and parses its content."""
        root = self._filesystem.inspect_get_roots()[0]
        inode = self._filesystem.stat('C:\\$Extend\\$UsnJrnl')['ino']

        with NamedTemporaryFile(buffering=0) as tempfile:
            self._filesystem.download_inode(root, inode, tempfile.name)

            journal = usn_journal(tempfile.name)

            return parse_journal(journal)

def parse_journal(journal):
    """Parses the USN Journal content removing duplicates
    and corrupted records.

    """
    events = [e for e in journal if not isinstance(e, CorruptedUsnRecord)]
    keyfunc = lambda e: str(e.file_reference_number) + e.file_name + e.timestamp
    event_groups = (tuple(g) for k, g in groupby(events, key=keyfunc))

    if len(events) < len(list(journal)):
        LOGGER.debug(
            "Corrupted records in UsnJrnl, some events might be missing.")

    return [journal_event(g) for g in event_groups]

def journal_event(events):
    """Group multiple events into a single one."""
    reasons = set(chain.from_iterable(e.reasons for e in events))
    attributes = set(chain.from_iterable(e.file_attributes for e in events))

    return JrnlEvent(events[0].file_reference_number,
                     events[0].parent_file_reference_number,
                     events[0].file_name,
                     events[0].timestamp,
                     list(reasons), list(attributes))

def generate_timeline(usnjrnl, content):
    """Aggregates the data collected from the USN journal
    and the filesystem content.

    """
    for event in usnjrnl:
        try:
            dirent = lookup_dirent(event, content)

            yield UsnJrnlEvent(
                dirent.inode, dirent.path, dirent.size, dirent.allocated,
                event.timestamp, event.changes, event.attributes)
        except LookupError as error:
            LOGGER.debug(error)

def lookup_dirent(event, content):
    for dirent in content[event.inode]:
        if dirent.path.endswith(event.name):
            return dirent

    # try constructing the full path from the parent folder
    for dirent in content[event.parent_inode]:
        if dirent.type == 'd':
            return Dirent(event.inode, ntpath.join(dirent.path, event.name), -1,
                          None, False, 0, 0, 0, 0)
    else:
        raise LookupError("File %s not found" % event.name)

def timestamp(secs, nsecs):
    delta = timedelta(seconds=secs) + timedelta(microseconds=(nsecs / 1000))

    return delta.total_seconds()

TSK_ALLOC = 0x01

Event = namedtuple('Event',
                   ('inode', 'path', 'size',
                    'allocated', 'timestamp', 'reason'))
UsnJrnlEvent = namedtuple('Event', ('file_reference_number', 'path', 'size',
                                    'allocated', 'timestamp', 'changes',
                                    'attributes'))

Dirent = namedtuple('Dirent', ('inode', 'path', 'size', 'type', 'allocated',
                               'atime', 'mtime', 'ctime', 'crtime'))
JrnlEvent = namedtuple('JrnlEvent', ('inode', 'parent_inode', 'name',
                                     'timestamp', 'changes', 'attributes'))

LOGGER = logging.getLogger("%s" % (__name__))
<commit_msg>usnjrnl_timeline: improve event lookup logic<commit_after># Copyright (c) 2016-2017, Matteo Cafasso
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.

# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS
# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
# OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
# OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Analyse disk content to extract File System event timelines."""

import ntpath
import logging
from datetime import timedelta
from functools import lru_cache
from itertools import chain, groupby
from tempfile import NamedTemporaryFile
from collections import defaultdict, namedtuple

from vminspect.filesystem import FileSystem
from vminspect.usnjrnl import CorruptedUsnRecord, usn_journal

class FSTimeline:
    def __init__(self, disk):
        self._disk = disk
        self._filesystem = None
        self._filetype_cache = {}
        self._checksum_cache = {}
        self.logger = logging.getLogger(
            "%s.%s" % (self.__module__, self.__class__.__name__))

    def __enter__(self):
        self._filesystem = FileSystem(self._disk)
        self._filesystem.mount()

        return self

    def __exit__(self, *_):
        self._filesystem.umount()

    def __getattr__(self, attr):
        return getattr(self._filesystem, attr)

    def timeline(self):
        self.logger.debug("Extracting File System timeline events.")
        events = tuple(Event(d.inode, d.path, d.size, d.allocated, t, r)
                       for d in self._visit_filesystem()
                       for t, r in ((d.atime, 'access'),
                                    (d.mtime, 'change'),
                                    (d.ctime, 'attribute_change'),
                                    (d.crtime, 'creation'))
                       if t > 0)

        self.logger.debug("Sorting File System timeline events.")
        return sorted(events, key=lambda e: e.timestamp)

    @lru_cache(maxsize=None)
    def file(self, path):
        """Identifies the file type.

        Caches the result to reduce overhead on duplicated events.

        """
        return self._filesystem.file(path)

    @lru_cache(maxsize=None)
    def checksum(self, path):
        """Identifies the file type.

        Caches the result to reduce overhead on duplicated events.

        """
        return self._filesystem.checksum(path)

    def _visit_filesystem(self):
        """Walks through the filesystem content."""
        self.logger.debug("Parsing File System content.")

        root_partition = self._filesystem.inspect_get_roots()[0]

        yield from self._root_dirent()

        for entry in self._filesystem.filesystem_walk(root_partition):
            yield Dirent(
                entry['tsk_inode'],
                self._filesystem.path('/' + entry['tsk_name']),
                entry['tsk_size'], entry['tsk_type'],
                True if entry['tsk_flags'] & TSK_ALLOC else False,
                timestamp(entry['tsk_atime_sec'], entry['tsk_atime_nsec']),
                timestamp(entry['tsk_mtime_sec'], entry['tsk_mtime_nsec']),
                timestamp(entry['tsk_ctime_sec'], entry['tsk_ctime_nsec']),
                timestamp(entry['tsk_crtime_sec'], entry['tsk_crtime_nsec']))

    def _root_dirent(self):
        """Returns the root folder dirent as filesystem_walk API doesn't."""
        fstat = self._filesystem.stat('/')

        yield Dirent(fstat['ino'], self._filesystem.path('/'),
                     fstat['size'], 'd', True,
                     timestamp(fstat['atime'], 0),
                     timestamp(fstat['mtime'], 0),
                     timestamp(fstat['ctime'], 0),
                     0)

class NTFSTimeline(FSTimeline):
    """Inspect NTFS filesystem in order to extract a timeline of events
    containing the information related to files/directories changes.

    This feature depends on a special build of Libguestfs available at:
      https://github.com/noxdafox/libguestfs/tree/forensics

    """
    def __init__(self, disk):
        super().__init__(disk)

    def __enter__(self):
        super().__enter__()

        if self._filesystem.osname != 'windows':
            self._filesystem.umount()

            raise RuntimeError("Filesystem is not NTFS.")

        return self

    def usnjrnl_timeline(self):
        """Iterates over the changes occurred within the filesystem.

        Yields UsnJrnlEvent namedtuples containing:

            file_reference_number: known in Unix FS as inode.
            path: full path of the file.
            size: size of the file in bytes if recoverable.
            allocated: whether the file exists or it has been deleted.
            timestamp: timespamp of the change.
            changes: list of changes applied to the file.
            attributes: list of file attributes.

        """
        filesystem_content = defaultdict(list)

        self.logger.debug("Extracting Update Sequence Number journal.")

        journal = self._read_journal()

        for dirent in self._visit_filesystem():
            filesystem_content[dirent.inode].append(dirent)

        self.logger.debug("Generating timeline.")
        yield from generate_timeline(journal, filesystem_content)

    def _read_journal(self):
        """Extracts the USN journal from the disk and parses its content."""
        root = self._filesystem.inspect_get_roots()[0]
        inode = self._filesystem.stat('C:\\$Extend\\$UsnJrnl')['ino']

        with NamedTemporaryFile(buffering=0) as tempfile:
            self._filesystem.download_inode(root, inode, tempfile.name)

            journal = usn_journal(tempfile.name)

            return parse_journal(journal)

def parse_journal(journal):
    """Parses the USN Journal content removing duplicates
    and corrupted records.

    """
    events = [e for e in journal if not isinstance(e, CorruptedUsnRecord)]
    keyfunc = lambda e: str(e.file_reference_number) + e.file_name + e.timestamp
    event_groups = (tuple(g) for k, g in groupby(events, key=keyfunc))

    if len(events) < len(list(journal)):
        LOGGER.debug(
            "Corrupted records in UsnJrnl, some events might be missing.")

    return [journal_event(g) for g in event_groups]

def journal_event(events):
    """Group multiple events into a single one."""
    reasons = set(chain.from_iterable(e.reasons for e in events))
    attributes = set(chain.from_iterable(e.file_attributes for e in events))

    return JrnlEvent(events[0].file_reference_number,
                     events[0].parent_file_reference_number,
                     events[0].file_name,
                     events[0].timestamp,
                     list(reasons), list(attributes))

def generate_timeline(usnjrnl, filesystem_content):
    """Aggregates the data collected from the USN journal
    and the filesystem content.

    """
    journal_content = defaultdict(list)
    for event in usnjrnl:
        journal_content[event.inode].append(event)

    for event in usnjrnl:
        try:
            dirent = lookup_dirent(event, filesystem_content, journal_content)

            yield UsnJrnlEvent(
                dirent.inode, dirent.path, dirent.size, dirent.allocated,
                event.timestamp, event.changes, event.attributes)
        except LookupError as error:
            LOGGER.debug(error)

def lookup_dirent(event, filesystem_content, journal_content):
    """Lookup the dirent given a journal event."""
    for dirent in filesystem_content[event.inode]:
        if dirent.path.endswith(event.name):
            return dirent

    path = lookup_folder(event, filesystem_content)
    if path is not None:
        return Dirent(event.inode, path, -1, None, False, 0, 0, 0, 0)

    path = lookup_deleted_folder(event, filesystem_content, journal_content)
    if path is not None:
        return Dirent(event.inode, path, -1, None, False, 0, 0, 0, 0)

    raise LookupError("File %s not found" % event.name)

def lookup_folder(event, filesystem):
    """Lookup the parent folder in the filesystem content."""
    for dirent in filesystem[event.parent_inode]:
        if dirent.type == 'd' and dirent.allocated:
            return ntpath.join(dirent.path, event.name)

def lookup_deleted_folder(event, filesystem, journal):
    """Lookup the parent folder in the journal content."""
    folder_events = (e for e in journal[event.parent_inode]
                     if 'DIRECTORY' in e.attributes
                     and 'FILE_DELETE' in e.changes)

    for folder_event in folder_events:
        path = lookup_deleted_folder(folder_event, filesystem, journal)

        return ntpath.join(path, event.name)

    return lookup_folder(event, filesystem)

def timestamp(secs, nsecs):
    delta = timedelta(seconds=secs) + timedelta(microseconds=(nsecs / 1000))

    return delta.total_seconds()

TSK_ALLOC = 0x01

Event = namedtuple('Event',
                   ('inode', 'path', 'size',
                    'allocated', 'timestamp', 'reason'))
UsnJrnlEvent = namedtuple('Event', ('file_reference_number', 'path', 'size',
                                    'allocated', 'timestamp', 'changes',
                                    'attributes'))

Dirent = namedtuple('Dirent', ('inode', 'path', 'size', 'type', 'allocated',
                               'atime', 'mtime', 'ctime', 'crtime'))
JrnlEvent = namedtuple('JrnlEvent', ('inode', 'parent_inode', 'name',
                                     'timestamp', 'changes', 'attributes'))

LOGGER = logging.getLogger("%s" % (__name__))
<|endoftext|>


<commit_before>"""
Adapted from https://github.com/numba/conda-recipe-cudatoolkit/blob/master/scripts/build.py

BSD 2-Clause License

Copyright (c) 2018 Onwards, Quansight, LLC
Copyright (c) 2017, Continuum Analytics, Inc.

All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """

from __future__ import print_function
import os
import sys
import shutil
import urllib.parse as urlparse
from pathlib import Path
import subprocess
from conda.exports import download, hashsum_file
import stat
import json

def set_chmod(file_name):
    # Do a simple chmod +x for a file within python
    st = os.stat(file_name)
    os.chmod(file_name, st.st_mode | stat.S_IXOTH)

def create_dir(new_dir):
    try:
        os.makedirs(new_dir)
    except FileExistsError:
        pass

def copy_files(src, dst):
    try:
        if (os.path.isfile(src)):
            set_chmod(src)
            shutil.copy(src, dst)
    except FileExistsError:
        pass

class Extractor(object):
    """Extractor base class, platform specific extractors should inherit
    from this class.
    """
    libdir = {'linux': 'lib',
              'osx': 'lib',
              'windows': 'DLLs', }

    def __init__(self, cudatoolkit_config, platform_config):
        """Initialise an instance:
        Arguments:
          cudatoolkit_config: the configuration for CUDA
          platform_config - the configuration for this platform
        """
        self.cu_name = cudatoolkit_config['name']
        self.cu_version = cudatoolkit_config['version']
        self.md5_url = cudatoolkit_config["md5_url"]
        self.base_url = cudatoolkit_config["base_url"]
        self.patch_url_text = cudatoolkit_config["patch_url_ext"]
        self.installers_url_ext = cudatoolkit_config["installers_url_ext"]
        self.cu_blob = platform_config['blob']
        self.conda_prefix = os.environ.get('CONDA_PREFIX')
        self.prefix = os.environ["PREFIX"]
        self.src_dir = Path(self.conda_prefix) / 'pkgs' / self.cu_name
        try:
            os.makedirs(self.src_dir)

        except FileExistsError:
            pass

        self.symlinks = getplatform() == "linux"
        self.debug_install_path = os.environ.get('DEBUG_INSTALLER_PATH')

    def create_activate_and_deactivate_scripts(self):
        activate_dir_path = Path(self.conda_prefix) / \
            'etc' / 'conda' / 'activate.d'
        deactivate_dir_path = Path(
            self.conda_prefix) / 'etc' / 'conda' / 'deactivate.d'

        try:
            os.makedirs(activate_dir_path)
            os.makedirs(deactivate_dir_path)

        except FileExistsError:
            pass

        # Copy cudatoolkit-dev-activate and cudatoolkit-dev-deactivate
        # to activate.d and deactivate.d directories

        scripts_dir = Path(self.prefix) / 'scripts'
        activate_scripts_dir = scripts_dir / 'activate.d'
        deactivate_scripts_dir = scripts_dir / 'deactivate.d'

        activate_scripts_list = [
            "cudatoolkit-dev-activate.sh",
            "cudatoolkit-dev-activate.bat"]
        for file_name in activate_scripts_list:
            file_full_path = activate_scripts_dir / file_name
            shutil.copy(file_full_path, activate_dir_path)

        deactivate_scripts_list = [
            "cudatoolkit-dev-deactivate.sh",
            "cudatoolkit-dev-deactivate.bat"]

        for file_name in deactivate_scripts_list:
            file_full_path = deactivate_scripts_dir / file_name
            shutil.copy(file_full_path, deactivate_dir_path)

    def download_blobs(self):
        """Downloads the binary blobs to the $SRC_DIR
        """
        dl_url = urlparse.urljoin(self.base_url, self.installers_url_ext)
        dl_url = urlparse.urljoin(dl_url, self.cu_blob)
        dl_path = os.path.join(self.src_dir, self.cu_blob)
        if not self.debug_install_path:
            print("downloading %s to %s" % (dl_url, dl_path))
            download(dl_url, dl_path)

        else:
            existing_file = os.path.join(self.debug_install_path, self.cu_blob)
            print("DEBUG: copying %s to %s" % (existing_file, dl_path))
            shutil.copy(existing_file, dl_path)

    def check_md5(self):
        """Checks the md5sums of the downloaded binaries
        """
        md5file = self.md5_url.split("/")[-1]
        path = os.path.join(self.src_dir, md5file)
        download(self.md5_url, path)

        # compute hash of blob
        blob_path = os.path.join(self.src_dir, self.cu_blob)
        md5sum = hashsum_file(blob_path, 'md5')

        # get checksums
        with open(path, 'r') as f:
            checksums = [x.strip().split() for x in f.read().splitlines() if x]

        # check md5 and filename match up
        check_dict = {x[0]: x[1] for x in checksums}
        assert check_dict[md5sum].startswith(self.cu_blob[:-7])

    def extract(self, *args):
        """The method to extract files from the cuda binary blobs.
        Platform specific extractors must implement.
        """
        raise RuntimeError("Must implement")

    def cleanup(self):
        """The method to delete unnecessary files after
        the installation process.
        """
        raise RuntimeError("Must implement")

class WindowsExtractor(Extractor):
    """The windows extractor
    """

    def extract(self):
        print("Extracting on Windows.....")
        runfile = os.path.join(self.src_dir, self.cu_blob)
        cmd = ['7za', 'x', '-o%s' %
               str(self.src_dir), runfile]
        try:
            subprocess.check_call(cmd)
        except subprocess.CalledProcessError as e:
            print("ERROR: Couldn't install Cudatoolkit: \
                   {reason}".format(reason=e))

    def cleanup(self):
        pass

class LinuxExtractor(Extractor):
    """The Linux Extractor
    """

    def extract(self):
        print("Extracting on Linux")
        runfile = os.path.join(self.src_dir, self.cu_blob)
        os.chmod(runfile, 0o777)
        cmd = [runfile, '--silent', '--toolkit',
               '--toolkitpath', str(self.src_dir), '--override']
        try:
            subprocess.check_call(cmd)
        except subprocess.CalledProcessError as e:
            print("ERROR: Couldn't install Cudatoolkit: \
                   {reason}".format(reason=e))

    def cleanup(self):
        blob_path = os.path.join(self.src_dir, self.cu_blob)
        if os.path.exists(blob_path):
            os.remove(blob_path)

        else:
            pass

class OsxExtractor(Extractor):
    """The osx Extractor
    """

    def _hdiutil_mount(self, temp_dir, file_name, install_dir):
        """Function to mount osx dmg images, extracts the files
           from an image into store and ensure they are
           unmounted on exit.
        """
        # open
        cmd = ['hdiutil', 'attach', '-mountpoint', temp_dir, file_name]
        cmd = ' '.join(cmd)
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
        process.wait()
        # find tar.gz files
        cmd = [
            'find',
            temp_dir,
            '-name',
            '"*.tar.gz"',
            '-exec',
            'tar',
            'xvf',
            "'{}'",
            '--directory',
            install_dir,
            "';'"]
        cmd = ' '.join(cmd)
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
        process.wait()
        # close
        cmd = ['hdiutil', 'detach', temp_dir]
        cmd = ' '.join(cmd)
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
        process.wait()

    def copy_files(self):
        src = Path(self.extract_temp_dir) / 'Developer' / \
            'NVIDIA' / 'CUDA-{}'.format(self.cu_version)
        dst = self.src_dir

        for f in os.listdir(src):
            source = src / f
            destination = dst / f
            try:
                shutil.copytree(source, destination)
            except NotADirectoryError:
                shutil.copy(source, destination)

    def extract(self):
        runfile = os.path.join(self.src_dir, self.cu_blob)
        extract_store_name = 'tmpstore'
        extract_temp_dir_name = 'tmp'
        self.extract_store = os.path.join(self.src_dir, extract_store_name)
        self.extract_temp_dir = os.path.join(
            self.src_dir, extract_temp_dir_name)
        create_dir(self.extract_store)
        create_dir(self.extract_temp_dir)
        self._hdiutil_mount(self.extract_store, runfile, self.extract_temp_dir)
        self.copy_files()

    def cleanup(self):
        blob_path = os.path.join(self.src_dir, self.cu_blob)
        if os.path.exists(blob_path):
            os.remove(blob_path)

        else:
            pass

        try:
            shutil.rmtree(self.extract_store)
            shutil.rmtree(self.extract_temp_dir)

        except FileNotFoundError:
            pass

def getplatform():
    plt = sys.platform
    if plt.startswith("linux"):
        return "linux"
    elif plt.startswith("win"):
        return "windows"
    elif plt.startswith("darwin"):
        return "osx"
    else:
        raise RuntimeError("Unknown platform")

def set_config():
    """Set necessary configurations"""

    cudatoolkit = {"linux": {}, "windows": {}, "osx": {}}
    prefix = Path(os.environ["PREFIX"])
    extra_args = dict()
    with open(prefix / 'bin' / 'cudatoolkit-dev-extra-args.txt', 'r') as f:
        extra_args = json.loads(f.read())

    # package version decl must match cuda release version
    cudatoolkit["version"] = os.environ['PKG_VERSION']
    cudatoolkit["name"] = os.environ['PKG_NAME']
    cudatoolkit["buildnum"] = os.environ['PKG_BUILDNUM']
    cudatoolkit["version_build"] = extra_args['version_build']
    cudatoolkit["driver_version"] = extra_args['driver_version']
    cudatoolkit["base_url"] = f'https://developer.nvidia.com/compute/cuda/{cudatoolkit["version"]}/Prod2/'
    cudatoolkit["installers_url_ext"] = f'local_installers/'
    cudatoolkit["patch_url_ext"] = f""
    cudatoolkit["md5_url"] = f'http://developer.download.nvidia.com/compute/cuda/{cudatoolkit["version"]}/Prod2/docs/sidebar/md5sum.txt'

    cudatoolkit["linux"] = {
        'blob': f'cuda_{cudatoolkit["version"]}.{cudatoolkit["version_build"]}_{cudatoolkit["driver_version"]}_linux',
    }

    cudatoolkit["windows"] = {
        'blob': f'cuda_{cudatoolkit["version"]}.{cudatoolkit["version_build"]}_windows', }

    cudatoolkit["osx"] = {
        'blob': f'cuda_{cudatoolkit["version"]}.{cudatoolkit["version_build"]}_mac', }

    return cudatoolkit

dispatcher = {
    "linux": LinuxExtractor,
    "windows": WindowsExtractor,
    "osx": OsxExtractor, }

def _main():

    print("Running Post installation")

    cudatoolkit_config = set_config()

    # get an extractor
    plat = getplatform()
    extractor_impl = dispatcher[plat]
    extractor = extractor_impl(
        cudatoolkit_config,
        cudatoolkit_config[plat])

    # create activate and deactivate scripts
    extractor.create_activate_and_deactivate_scripts()

    # download binaries
    extractor.download_blobs()

    # check md5sum
    extractor.check_md5()

    # Extract
    extractor.extract()

    # Cleanup
    extractor.cleanup()

if __name__ == "__main__":
    _main()
<commit_msg>add FileExistsError exception<commit_after>"""
Adapted from https://github.com/numba/conda-recipe-cudatoolkit/blob/master/scripts/build.py

BSD 2-Clause License

Copyright (c) 2018 Onwards, Quansight, LLC
Copyright (c) 2017, Continuum Analytics, Inc.

All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """

from __future__ import print_function
import os
import sys
import shutil
import urllib.parse as urlparse
from pathlib import Path
import subprocess
from conda.exports import download, hashsum_file
import stat
import json

def set_chmod(file_name):
    # Do a simple chmod +x for a file within python
    st = os.stat(file_name)
    os.chmod(file_name, st.st_mode | stat.S_IXOTH)

def create_dir(new_dir):
    try:
        os.makedirs(new_dir)
    except FileExistsError:
        pass

def copy_files(src, dst):
    try:
        if (os.path.isfile(src)):
            set_chmod(src)
            shutil.copy(src, dst)
    except FileExistsError:
        pass

class Extractor(object):
    """Extractor base class, platform specific extractors should inherit
    from this class.
    """
    libdir = {'linux': 'lib',
              'osx': 'lib',
              'windows': 'DLLs', }

    def __init__(self, cudatoolkit_config, platform_config):
        """Initialise an instance:
        Arguments:
          cudatoolkit_config: the configuration for CUDA
          platform_config - the configuration for this platform
        """
        self.cu_name = cudatoolkit_config['name']
        self.cu_version = cudatoolkit_config['version']
        self.md5_url = cudatoolkit_config["md5_url"]
        self.base_url = cudatoolkit_config["base_url"]
        self.patch_url_text = cudatoolkit_config["patch_url_ext"]
        self.installers_url_ext = cudatoolkit_config["installers_url_ext"]
        self.cu_blob = platform_config['blob']
        self.conda_prefix = os.environ.get('CONDA_PREFIX')
        self.prefix = os.environ["PREFIX"]
        self.src_dir = Path(self.conda_prefix) / 'pkgs' / self.cu_name
        try:
            os.makedirs(self.src_dir)

        except FileExistsError:
            pass

        self.symlinks = getplatform() == "linux"
        self.debug_install_path = os.environ.get('DEBUG_INSTALLER_PATH')

    def create_activate_and_deactivate_scripts(self):
        activate_dir_path = Path(self.conda_prefix) / \
            'etc' / 'conda' / 'activate.d'
        deactivate_dir_path = Path(
            self.conda_prefix) / 'etc' / 'conda' / 'deactivate.d'

        try:
            os.makedirs(activate_dir_path)
            os.makedirs(deactivate_dir_path)

        except FileExistsError:
            pass

        # Copy cudatoolkit-dev-activate and cudatoolkit-dev-deactivate
        # to activate.d and deactivate.d directories

        scripts_dir = Path(self.prefix) / 'scripts'
        activate_scripts_dir = scripts_dir / 'activate.d'
        deactivate_scripts_dir = scripts_dir / 'deactivate.d'

        activate_scripts_list = [
            "cudatoolkit-dev-activate.sh",
            "cudatoolkit-dev-activate.bat"]
        for file_name in activate_scripts_list:
            file_full_path = activate_scripts_dir / file_name
            shutil.copy(file_full_path, activate_dir_path)

        deactivate_scripts_list = [
            "cudatoolkit-dev-deactivate.sh",
            "cudatoolkit-dev-deactivate.bat"]

        for file_name in deactivate_scripts_list:
            file_full_path = deactivate_scripts_dir / file_name
            shutil.copy(file_full_path, deactivate_dir_path)

    def download_blobs(self):
        """Downloads the binary blobs to the $SRC_DIR
        """
        dl_url = urlparse.urljoin(self.base_url, self.installers_url_ext)
        dl_url = urlparse.urljoin(dl_url, self.cu_blob)
        dl_path = os.path.join(self.src_dir, self.cu_blob)
        if not self.debug_install_path:
            print("downloading %s to %s" % (dl_url, dl_path))
            download(dl_url, dl_path)

        else:
            existing_file = os.path.join(self.debug_install_path, self.cu_blob)
            print("DEBUG: copying %s to %s" % (existing_file, dl_path))
            shutil.copy(existing_file, dl_path)

    def check_md5(self):
        """Checks the md5sums of the downloaded binaries
        """
        md5file = self.md5_url.split("/")[-1]
        path = os.path.join(self.src_dir, md5file)
        download(self.md5_url, path)

        # compute hash of blob
        blob_path = os.path.join(self.src_dir, self.cu_blob)
        md5sum = hashsum_file(blob_path, 'md5')

        # get checksums
        with open(path, 'r') as f:
            checksums = [x.strip().split() for x in f.read().splitlines() if x]

        # check md5 and filename match up
        check_dict = {x[0]: x[1] for x in checksums}
        assert check_dict[md5sum].startswith(self.cu_blob[:-7])

    def extract(self, *args):
        """The method to extract files from the cuda binary blobs.
        Platform specific extractors must implement.
        """
        raise RuntimeError("Must implement")

    def cleanup(self):
        """The method to delete unnecessary files after
        the installation process.
        """
        raise RuntimeError("Must implement")

class WindowsExtractor(Extractor):
    """The windows extractor
    """

    def extract(self):
        print("Extracting on Windows.....")
        runfile = os.path.join(self.src_dir, self.cu_blob)
        cmd = ['7za', 'x', '-o%s' %
               str(self.src_dir), runfile]
        try:
            subprocess.check_call(cmd)
        except subprocess.CalledProcessError as e:
            print("ERROR: Couldn't install Cudatoolkit: \
                   {reason}".format(reason=e))

    def cleanup(self):
        pass

class LinuxExtractor(Extractor):
    """The Linux Extractor
    """

    def extract(self):
        print("Extracting on Linux")
        runfile = os.path.join(self.src_dir, self.cu_blob)
        os.chmod(runfile, 0o777)
        cmd = [runfile, '--silent', '--toolkit',
               '--toolkitpath', str(self.src_dir), '--override']
        try:
            subprocess.check_call(cmd)
        except subprocess.CalledProcessError as e:
            print("ERROR: Couldn't install Cudatoolkit: \
                   {reason}".format(reason=e))

    def cleanup(self):
        blob_path = os.path.join(self.src_dir, self.cu_blob)
        if os.path.exists(blob_path):
            os.remove(blob_path)

        else:
            pass

class OsxExtractor(Extractor):
    """The osx Extractor
    """

    def _hdiutil_mount(self, temp_dir, file_name, install_dir):
        """Function to mount osx dmg images, extracts the files
           from an image into store and ensure they are
           unmounted on exit.
        """
        # open
        cmd = ['hdiutil', 'attach', '-mountpoint', temp_dir, file_name]
        cmd = ' '.join(cmd)
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
        process.wait()
        # find tar.gz files
        cmd = [
            'find',
            temp_dir,
            '-name',
            '"*.tar.gz"',
            '-exec',
            'tar',
            'xvf',
            "'{}'",
            '--directory',
            install_dir,
            "';'"]
        cmd = ' '.join(cmd)
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
        process.wait()
        # close
        cmd = ['hdiutil', 'detach', temp_dir]
        cmd = ' '.join(cmd)
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
        process.wait()

    def copy_files(self):
        src = Path(self.extract_temp_dir) / 'Developer' / \
            'NVIDIA' / 'CUDA-{}'.format(self.cu_version)
        dst = self.src_dir

        for f in os.listdir(src):
            source = src / f
            destination = dst / f
            try:
                shutil.copytree(source, destination)
            except NotADirectoryError:
                shutil.copy(source, destination)
            except FileExistsError:
                pass

    def extract(self):
        runfile = os.path.join(self.src_dir, self.cu_blob)
        extract_store_name = 'tmpstore'
        extract_temp_dir_name = 'tmp'
        self.extract_store = os.path.join(self.src_dir, extract_store_name)
        self.extract_temp_dir = os.path.join(
            self.src_dir, extract_temp_dir_name)
        create_dir(self.extract_store)
        create_dir(self.extract_temp_dir)
        self._hdiutil_mount(self.extract_store, runfile, self.extract_temp_dir)
        self.copy_files()

    def cleanup(self):
        blob_path = os.path.join(self.src_dir, self.cu_blob)
        if os.path.exists(blob_path):
            os.remove(blob_path)

        else:
            pass

        try:
            shutil.rmtree(self.extract_store)
            shutil.rmtree(self.extract_temp_dir)

        except FileNotFoundError:
            pass

def getplatform():
    plt = sys.platform
    if plt.startswith("linux"):
        return "linux"
    elif plt.startswith("win"):
        return "windows"
    elif plt.startswith("darwin"):
        return "osx"
    else:
        raise RuntimeError("Unknown platform")

def set_config():
    """Set necessary configurations"""

    cudatoolkit = {"linux": {}, "windows": {}, "osx": {}}
    prefix = Path(os.environ["PREFIX"])
    extra_args = dict()
    with open(prefix / 'bin' / 'cudatoolkit-dev-extra-args.txt', 'r') as f:
        extra_args = json.loads(f.read())

    # package version decl must match cuda release version
    cudatoolkit["version"] = os.environ['PKG_VERSION']
    cudatoolkit["name"] = os.environ['PKG_NAME']
    cudatoolkit["buildnum"] = os.environ['PKG_BUILDNUM']
    cudatoolkit["version_build"] = extra_args['version_build']
    cudatoolkit["driver_version"] = extra_args['driver_version']
    cudatoolkit["base_url"] = f'https://developer.nvidia.com/compute/cuda/{cudatoolkit["version"]}/Prod2/'
    cudatoolkit["installers_url_ext"] = f'local_installers/'
    cudatoolkit["patch_url_ext"] = f""
    cudatoolkit["md5_url"] = f'http://developer.download.nvidia.com/compute/cuda/{cudatoolkit["version"]}/Prod2/docs/sidebar/md5sum.txt'

    cudatoolkit["linux"] = {
        'blob': f'cuda_{cudatoolkit["version"]}.{cudatoolkit["version_build"]}_{cudatoolkit["driver_version"]}_linux',
    }

    cudatoolkit["windows"] = {
        'blob': f'cuda_{cudatoolkit["version"]}.{cudatoolkit["version_build"]}_windows', }

    cudatoolkit["osx"] = {
        'blob': f'cuda_{cudatoolkit["version"]}.{cudatoolkit["version_build"]}_mac', }

    return cudatoolkit

dispatcher = {
    "linux": LinuxExtractor,
    "windows": WindowsExtractor,
    "osx": OsxExtractor, }

def _main():

    print("Running Post installation")

    cudatoolkit_config = set_config()

    # get an extractor
    plat = getplatform()
    extractor_impl = dispatcher[plat]
    extractor = extractor_impl(
        cudatoolkit_config,
        cudatoolkit_config[plat])

    # create activate and deactivate scripts
    extractor.create_activate_and_deactivate_scripts()

    # download binaries
    extractor.download_blobs()

    # check md5sum
    extractor.check_md5()

    # Extract
    extractor.extract()

    # Cleanup
    extractor.cleanup()

if __name__ == "__main__":
    _main()
<|endoftext|>


<commit_before>from trivup import trivup
from trivup.apps.KerberosKdcApp import KerberosKdcApp

import os
import time

class KafkaBrokerApp (trivup.App):
    """ Kafka broker app
        Depends on ZookeeperApp """
    def __init__(self, cluster, conf=None, on=None):
        """
        @param cluster     Current cluster
        @param conf        Configuration dict, see below.
        @param on          Node name to run on

        Supported conf keys:
           * version - Kafka version to use, will build 'trunk' from kafka_path,
                       otherwise the version is taken to be a formal release which
                       will be downloaded and deployed.
           * listeners - CSV list of listener types: PLAINTEXT,SSL,SASL,SASL_SSL
           * listener_host - alternative listener host instead of node name (e.g., '*')
           * sasl_mechanisms - CSV list of SASL mechanisms to enable: GSSAPI,PLAIN
                               SASL listeners will be added automatically.
                               KerberosKdcApp is required for GSSAPI.
           * sasl_users - CSV list of SASL PLAIN of user=pass for authenticating clients
           * ssl_client_auth - ssl.client.auth broker property (def: required)
           * num_partitions - Topic auto-create partition count (3)
           * replication_Factor - Topic auto-create replication factor (1)
           * port_base - Low TCP port base to start allocating from (random)
           * kafka_path - Path to Kafka build tree (for trunk usage)
           * conf - arbitary server.properties config as a list of strings.
        """
        super(KafkaBrokerApp, self).__init__(cluster, conf=conf, on=on)

        self.zk = cluster.find_app('ZookeeperApp')
        if self.zk is None:
            raise Exception('ZookeeperApp required')

        # Kafka repo uses SVN nomenclature
        if self.conf['version'] == 'master':
            self.conf['version'] = 'trunk'

        listener_host = self.conf.get('listener_host', self.conf.get('nodename'))

        # Arbitrary (non-template) configuration statements
        conf_blob = self.conf.get('conf', list())
        jaas_blob = list()

        #
        # Configure listeners, SSL and SASL
        #
        listeners = self.conf.get('listeners', 'PLAINTEXT').split(',')

        # SASL support
        sasl_mechs = [x for x in self.conf.get('sasl_mechanisms', '').replace(' ', '').split(',') if len(x) > 0]
        if len(sasl_mechs) > 0:
            listeners.append('SASL_PLAINTEXT')

        # SSL support
        if getattr(cluster, 'ssl', None) is not None:
            # Add SSL listeners
            listeners.append('SSL')
            if len(sasl_mechs) > 0:
                listeners.append('SASL_SSL')

        # Create listeners
        ports = [(x, trivup.TcpPortAllocator(self.cluster).next(self, self.conf.get('port_base', None))) for x in sorted(set(listeners))]
        self.conf['port'] = ports[0][1] # "Default" port
        self.conf['address'] = '%s:%d' % (listener_host, self.conf['port'])
        self.conf['listeners'] = ','.join(['%s://%s:%d' % (x[0], listener_host, x[1]) for x in ports])
        self.conf['advertised.listeners'] = self.conf['listeners']
        self.conf['auto_create_topics'] = self.conf.get('auto_create_topics', 'true')
        self.dbg('Listeners: %s' % self.conf['listeners'])
        self.dbg('Advertised Listeners: %s' % self.conf['advertised_listeners'])

        if len(sasl_mechs) > 0:
            self.dbg('SASL mechanisms: %s' % sasl_mechs)
            jaas_blob.append('KafkaServer {')

            conf_blob.append('sasl.enabled.mechanisms=%s' % ','.join(sasl_mechs))
            if 'PLAIN' in sasl_mechs:
                sasl_users = self.conf.get('sasl_users', '')
                if len(sasl_users) == 0:
                    self.log('WARNING: No sasl_users configured for PLAIN, expected CSV of user=pass,..')
                else:
                    jaas_blob.append('org.apache.kafka.common.security.plain.PlainLoginModule required debug=true')
                    for up in sasl_users.split(','):
                        u,p = up.split('=')
                        jaas_blob.append('user_%s="%s"' % (u, p))
                    jaas_blob[-1] += ';'

            if 'GSSAPI' in sasl_mechs:
                conf_blob.append('sasl.kerberos.service.name=%s' % 'kafka')
                kdc = self.cluster.find_app(KerberosKdcApp)
                self.env_add('KRB5_CONFIG', kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Djava.security.krb5.conf=%s' % kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Dsun.security.krb5.debug=true')
                self.kerberos_principal,self.kerberos_keytab = kdc.add_principal('kafka', listener_host)
                jaas_blob.append('com.sun.security.auth.module.Krb5LoginModule required')
                jaas_blob.append('useKeyTab=true storeKey=true doNotPrompt=true')
                jaas_blob.append('keyTab="%s"' % self.kerberos_keytab)
                jaas_blob.append('debug=true')
                jaas_blob.append('principal="%s";' % self.kerberos_principal)

            jaas_blob.append('};\n')
            self.conf['jaas_file'] = self.create_file('jaas_broker.conf', data='\n'.join(jaas_blob))
            self.env_add('KAFKA_OPTS', '-Djava.security.auth.login.config=%s' % self.conf['jaas_file'])
            self.env_add('KAFKA_OPTS', '-Djava.security.debug=all')

        # SSL config and keys (et.al.)
        if getattr(cluster, 'ssl', None) is not None:
            ssl = cluster.ssl
            keystore,truststore,_,_ = ssl.create_keystore('broker%s' % self.appid)
            conf_blob.append('ssl.protocol=TLS')
            conf_blob.append('ssl.enabled.protocols=TLSv1.2,TLSv1.1,TLSv1')
            conf_blob.append('ssl.keystore.type = JKS')
            conf_blob.append('ssl.keystore.location = %s' % keystore)
            conf_blob.append('ssl.keystore.password = %s ' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.key.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.truststore.type = JKS')
            conf_blob.append('ssl.truststore.location = %s' % truststore)
            conf_blob.append('ssl.truststore.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.client.auth = %s' % self.conf.get('ssl_client_auth', 'required'))

        # Kafka Configuration properties
        self.conf['log_dirs'] = self.create_dir('logs')
        if 'num_partitions' not in self.conf:
            self.conf['num_partitions'] = 3
        self.conf['zk_connect'] = self.zk.get('address', None)
        if 'replication_factor' not in self.conf:
            self.conf['replication_factor'] = 1

        # Generate config file
        self.conf['conf_file'] = self.create_file_from_template('server.properties',
                                                                self.conf,
                                                                append_data='\n'.join(conf_blob))

        # Generate LOG4J file (if app debug is enabled)
        if self.debug:
            self.conf['log4j_file'] = self.create_file_from_template('log4j.properties', self.conf, subst=False)
            self.env_add('KAFKA_LOG4J_OPTS', '-Dlog4j.configuration=file:%s' % self.conf['log4j_file'])

        # Runs in foreground, stopped by Ctrl-C
        # This is the default for no-deploy use: will be overwritten by deploy() if enabled.
        if self.conf.get('kafka_path', None) is not None:
            self.conf['destdir'] = self.conf['kafka_path']
            self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
            start_sh = os.path.join(self.conf['bindir'], 'kafka-server-start.sh')
        else:
            start_sh = 'kafka-server-start.sh'

        self.conf['start_cmd'] = '%s %s' % (start_sh, self.conf['conf_file'])
        self.conf['stop_cmd'] = None # Ctrl-C

    def operational (self):
        self.dbg('Checking if operational')
        return os.system('(echo anything | nc %s) 2>/dev/null' %
                         ' '.join(self.get('address').split(':'))) == 0

    def deploy (self):
        destdir = os.path.join(self.cluster.mkpath(self.__class__.__name__), 'kafka', self.get('version'))
        self.dbg('Deploy %s version %s on %s to %s' %
                 (self.name, self.get('version'), self.node.name, destdir))
        deploy_exec = self.resource_path('deploy.sh')
        if not os.path.exists(deploy_exec):
            raise NotImplementedError('Kafka deploy.sh script missing in %s' %
                                      deploy_exec)
        t_start = time.time()
        cmd = '%s %s "%s" "%s"' % \
              (deploy_exec, self.get('version'), self.get('kafka_path', destdir), destdir)
        r = os.system(cmd)
        if r != 0:
            raise Exception('Deploy "%s" returned exit code %d' % (cmd, r))
        self.dbg('Deployed version %s in %ds' %
                 (self.get('version'), time.time() - t_start))

        self.conf['destdir'] = destdir
        self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
        # Override start command with updated path.
        self.conf['start_cmd'] = '%s/bin/kafka-server-start.sh %s' % (destdir, self.conf['conf_file'])
        self.dbg('Updated start_cmd to %s' % self.conf['start_cmd'])
<commit_msg>Maybe prepare for SASL SCRAM<commit_after>from trivup import trivup
from trivup.apps.KerberosKdcApp import KerberosKdcApp

import os
import time

class KafkaBrokerApp (trivup.App):
    """ Kafka broker app
        Depends on ZookeeperApp """
    def __init__(self, cluster, conf=None, on=None):
        """
        @param cluster     Current cluster
        @param conf        Configuration dict, see below.
        @param on          Node name to run on

        Supported conf keys:
           * version - Kafka version to use, will build 'trunk' from kafka_path,
                       otherwise the version is taken to be a formal release which
                       will be downloaded and deployed.
           * listeners - CSV list of listener types: PLAINTEXT,SSL,SASL,SASL_SSL
           * listener_host - alternative listener host instead of node name (e.g., '*')
           * sasl_mechanisms - CSV list of SASL mechanisms to enable: GSSAPI,PLAIN
                               SASL listeners will be added automatically.
                               KerberosKdcApp is required for GSSAPI.
           * sasl_users - CSV list of SASL PLAIN of user=pass for authenticating clients
           * ssl_client_auth - ssl.client.auth broker property (def: required)
           * num_partitions - Topic auto-create partition count (3)
           * replication_Factor - Topic auto-create replication factor (1)
           * port_base - Low TCP port base to start allocating from (random)
           * kafka_path - Path to Kafka build tree (for trunk usage)
           * conf - arbitary server.properties config as a list of strings.
        """
        super(KafkaBrokerApp, self).__init__(cluster, conf=conf, on=on)

        self.zk = cluster.find_app('ZookeeperApp')
        if self.zk is None:
            raise Exception('ZookeeperApp required')

        # Kafka repo uses SVN nomenclature
        if self.conf['version'] == 'master':
            self.conf['version'] = 'trunk'

        listener_host = self.conf.get('listener_host', self.conf.get('nodename'))

        # Arbitrary (non-template) configuration statements
        conf_blob = self.conf.get('conf', list())
        jaas_blob = list()

        #
        # Configure listeners, SSL and SASL
        #
        listeners = self.conf.get('listeners', 'PLAINTEXT').split(',')

        # SASL support
        sasl_mechs = [x for x in self.conf.get('sasl_mechanisms', '').replace(' ', '').split(',') if len(x) > 0]
        if len(sasl_mechs) > 0:
            listeners.append('SASL_PLAINTEXT')

        # SSL support
        if getattr(cluster, 'ssl', None) is not None:
            # Add SSL listeners
            listeners.append('SSL')
            if len(sasl_mechs) > 0:
                listeners.append('SASL_SSL')

        # Create listeners
        ports = [(x, trivup.TcpPortAllocator(self.cluster).next(self, self.conf.get('port_base', None))) for x in sorted(set(listeners))]
        self.conf['port'] = ports[0][1] # "Default" port
        self.conf['address'] = '%s:%d' % (listener_host, self.conf['port'])
        self.conf['listeners'] = ','.join(['%s://%s:%d' % (x[0], listener_host, x[1]) for x in ports])
        self.conf['advertised.listeners'] = self.conf['listeners']
        self.conf['auto_create_topics'] = self.conf.get('auto_create_topics', 'true')
        self.dbg('Listeners: %s' % self.conf['listeners'])
        self.dbg('Advertised Listeners: %s' % self.conf['advertised_listeners'])

        if len(sasl_mechs) > 0:
            self.dbg('SASL mechanisms: %s' % sasl_mechs)
            jaas_blob.append('KafkaServer {')

            conf_blob.append('sasl.enabled.mechanisms=%s' % ','.join(sasl_mechs))
            if 'PLAIN' in sasl_mechs or 'SCRAM-SHA-1' in sasl_mechs:
                if 'PLAIN' in sasl_mechs:
                    plugin = 'Plain'
                else:
                    plugin = 'Scram'
                sasl_users = self.conf.get('sasl_users', '')
                if len(sasl_users) == 0:
                    self.log('WARNING: No sasl_users configured for %s, expected CSV of user=pass,..' % plugin)
                else:
                    jaas_blob.append('org.apache.kafka.common.security.plain.%sLoginModule required debug=true' % plugin)
                    for up in sasl_users.split(','):
                        u,p = up.split('=')
                        jaas_blob.append('user_%s="%s"' % (u, p))
                    jaas_blob[-1] += ';'

            if 'GSSAPI' in sasl_mechs:
                conf_blob.append('sasl.kerberos.service.name=%s' % 'kafka')
                kdc = self.cluster.find_app(KerberosKdcApp)
                self.env_add('KRB5_CONFIG', kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Djava.security.krb5.conf=%s' % kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Dsun.security.krb5.debug=true')
                self.kerberos_principal,self.kerberos_keytab = kdc.add_principal('kafka', listener_host)
                jaas_blob.append('com.sun.security.auth.module.Krb5LoginModule required')
                jaas_blob.append('useKeyTab=true storeKey=true doNotPrompt=true')
                jaas_blob.append('keyTab="%s"' % self.kerberos_keytab)
                jaas_blob.append('debug=true')
                jaas_blob.append('principal="%s";' % self.kerberos_principal)

            jaas_blob.append('};\n')
            self.conf['jaas_file'] = self.create_file('jaas_broker.conf', data='\n'.join(jaas_blob))
            self.env_add('KAFKA_OPTS', '-Djava.security.auth.login.config=%s' % self.conf['jaas_file'])
            self.env_add('KAFKA_OPTS', '-Djava.security.debug=all')

        # SSL config and keys (et.al.)
        if getattr(cluster, 'ssl', None) is not None:
            ssl = cluster.ssl
            keystore,truststore,_,_ = ssl.create_keystore('broker%s' % self.appid)
            conf_blob.append('ssl.protocol=TLS')
            conf_blob.append('ssl.enabled.protocols=TLSv1.2,TLSv1.1,TLSv1')
            conf_blob.append('ssl.keystore.type = JKS')
            conf_blob.append('ssl.keystore.location = %s' % keystore)
            conf_blob.append('ssl.keystore.password = %s ' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.key.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.truststore.type = JKS')
            conf_blob.append('ssl.truststore.location = %s' % truststore)
            conf_blob.append('ssl.truststore.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.client.auth = %s' % self.conf.get('ssl_client_auth', 'required'))

        # Kafka Configuration properties
        self.conf['log_dirs'] = self.create_dir('logs')
        if 'num_partitions' not in self.conf:
            self.conf['num_partitions'] = 3
        self.conf['zk_connect'] = self.zk.get('address', None)
        if 'replication_factor' not in self.conf:
            self.conf['replication_factor'] = 1

        # Generate config file
        self.conf['conf_file'] = self.create_file_from_template('server.properties',
                                                                self.conf,
                                                                append_data='\n'.join(conf_blob))

        # Generate LOG4J file (if app debug is enabled)
        if self.debug:
            self.conf['log4j_file'] = self.create_file_from_template('log4j.properties', self.conf, subst=False)
            self.env_add('KAFKA_LOG4J_OPTS', '-Dlog4j.configuration=file:%s' % self.conf['log4j_file'])

        # Runs in foreground, stopped by Ctrl-C
        # This is the default for no-deploy use: will be overwritten by deploy() if enabled.
        if self.conf.get('kafka_path', None) is not None:
            self.conf['destdir'] = self.conf['kafka_path']
            self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
            start_sh = os.path.join(self.conf['bindir'], 'kafka-server-start.sh')
        else:
            start_sh = 'kafka-server-start.sh'

        self.conf['start_cmd'] = '%s %s' % (start_sh, self.conf['conf_file'])
        self.conf['stop_cmd'] = None # Ctrl-C

    def operational (self):
        self.dbg('Checking if operational')
        return os.system('(echo anything | nc %s) 2>/dev/null' %
                         ' '.join(self.get('address').split(':'))) == 0

    def deploy (self):
        destdir = os.path.join(self.cluster.mkpath(self.__class__.__name__), 'kafka', self.get('version'))
        self.dbg('Deploy %s version %s on %s to %s' %
                 (self.name, self.get('version'), self.node.name, destdir))
        deploy_exec = self.resource_path('deploy.sh')
        if not os.path.exists(deploy_exec):
            raise NotImplementedError('Kafka deploy.sh script missing in %s' %
                                      deploy_exec)
        t_start = time.time()
        cmd = '%s %s "%s" "%s"' % \
              (deploy_exec, self.get('version'), self.get('kafka_path', destdir), destdir)
        r = os.system(cmd)
        if r != 0:
            raise Exception('Deploy "%s" returned exit code %d' % (cmd, r))
        self.dbg('Deployed version %s in %ds' %
                 (self.get('version'), time.time() - t_start))

        self.conf['destdir'] = destdir
        self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
        # Override start command with updated path.
        self.conf['start_cmd'] = '%s/bin/kafka-server-start.sh %s' % (destdir, self.conf['conf_file'])
        self.dbg('Updated start_cmd to %s' % self.conf['start_cmd'])
<|endoftext|>


<commit_before>import csv
import os
from os import walk

PATH = '/home/darth/Documents/Adamson University/CS Research Project/Kyoto 2013 Dataset/CSV'

data = []

def main():
	# list all CSV files under the PATH
	data = list_files(path=PATH)
	for datum in data:
		# read CSV data
		in_csv = csv.reader(open(datum, 'r'))

def list_files(path):
	file_list = []
	for (dirpath, dirnames, filenames) in walk(path):
		file_list.extend(os.path.join(dirpath, filename) for filename in filenames)
	return file_list

def feature_scale(x):
	'''Scales integer values [min, max] -> [0.0, 1.0]'''
	return [ (x_i - min(x)) / (max(x) - min(x)) for x_i in x ]

def scale_boolean_values(value):
	'''Scales the supposed boolean values in the dataset'''
	if (value != 0):
		return 1
	elif (value == 0):
		return 0

def scale_label(label):
	'''Scales the labels [-2 -1 1] -> [0, 1]'''
	if (label == -2 or label == -1):
		return 1	# -2 and -1 represents an attack
	elif (label == 1):
		return 0	# 1 represents normal session

def convert_time_to_integer(time):
	'''Converts time to integer value, preparation for scaling'''
	h, m, s = map(int, time.split(':'))
	parsed_time = h + (m * (1 / 60)) + (s * (1 / 3600))
	return parsed_time

def standardize(x):
	'''Method for normalizing using Student's t-statistic'''
	mean = sum(x) / len(x)
	std_dev = (1 / len(x) * sum([ (x_i - mean) ** 2 for x_i in x] )) ** 0.5
	z_score = [ (x_i - mean) / std_dev for x_i in x ]
	return z_score

if __name__ == '__main__':
	main()<commit_msg>standardize-data.py was the module used for preprocessing the dataset, other than txt-to-csv.py<commit_after><|endoftext|>