<commit_before>"""
github.py - Willie Github Module
Copyright 2012, Dimitri Molenaars http://tyrope.nl/
Licensed under the Eiffel Forum License 2.

http://willie.dftba.net/
"""

from datetime import datetime
from urllib2 import HTTPError
import json
from willie import web
from willie.module import commands
import os

def checkConfig(bot):
    if not bot.config.has_option('github', 'oauth_token') or not bot.config.has_option('github', 'repo'):
        return False
    else:
        return [bot.config.github.oauth_token, bot.config.github.repo]

def configure(config):
    """
    | [github] | example | purpose |
    | -------- | ------- | ------- |
    | oauth_token | 5868e7af57496cc3ae255868e7af57496cc3ae25 | The OAuth token to connect to your github repo |
    | repo | embolalia/willie | The GitHub repo you're working from. |
    """
    chunk = ''
    if config.option('Configuring github issue reporting and searching module', False):
        config.interactive_add('github', 'oauth_token', 'Github API Oauth2 token', '')
        config.interactive_add('github', 'repo', 'Github repository', 'embolalia/willie')
    return chunk

@commands('makeissue', 'makebug')
def issue(bot, trigger):
    """Create a GitHub issue, also known as a bug report. Syntax: .makeissue Title of the bug report"""
    # check input
    if not trigger.group(2):
        return bot.say('Please title the issue')

    # Is the Oauth token and repo available?
    gitAPI = checkConfig(bot)
    if not gitAPI:
        return bot.say('Git module not configured, make sure github.oauth_token and github.repo are defined')

    # parse input
    now = ' '.join(str(datetime.utcnow()).split(' ')).split('.')[0] + ' UTC'
    body = 'Submitted by: %s\nFrom channel: %s\nAt %s' % (trigger.nick, trigger.sender, now)
    data = {"title": trigger.group(2).encode('utf-8'), "body": body}
    # submit
    try:
        raw = web.post('https://api.github.com/repos/' + gitAPI[1] + '/issues?access_token=' + gitAPI[0], json.dumps(data))
    except HTTPError:
        return bot.say('The GitHub API returned an error.')

    data = json.loads(raw)
    bot.say('Issue #%s posted. %s' % (data['number'], data['html_url']))
    bot.debug(__file__, 'Issue #%s created in %s' % (data['number'], trigger.sender), 'warning')

@commands('addtrace', 'addtraceback')
def add_traceback(bot, trigger):
    """Add a traceback to a GitHub issue.

    This pulls the traceback from the exceptions log file. To use, put .addtrace
    followed by the issue number to add the comment to, then the signature of
    the error (the message shown to the channel when the error occured). This
    command will only work for errors from unhandled exceptions."""
    # Make sure the API is set up
    gitAPI = checkConfig(bot)
    if not gitAPI:
        return bot.say('Git module not configured, make sure github.oauth_token and github.repo are defined')

    # Make sure the input is valid
    args = trigger.group(2).split(None, 1)
    if len(args) != 2:
        bot.say(validate)
        return
    number, trace = args

    # Make sure the given issue number exists
    issue_data = web.get('https://api.github.com/repos/%s/issues/%s' % (gitAPI[1], number))
    issue_data = json.loads(issue_data)
    if 'message' in issue_data and issue_data['message'] == 'Not Found':
        return bot.say("That issue doesn't exist.")

    # Find the relevant lines from the log file
    post = ''
    logfile = os.path.join(bot.config.logdir, 'exceptions.log')
    with open(logfile) as log:
        in_trace = False
        for data in log:
            if data == 'Signature: ' + trace + '\n':
                post = data
                in_trace = True
            elif data == '----------------------------------------\n':
                in_trace = False
            elif in_trace:
                post += data

    # Give an error if we didn't find the traceback
    if not post:
        return bot.say("I don't remember getting that error. Please post it "
                       "yourself at https://github.com/%s/issues/%s"
                       % (gitAPI[1], number))

    # Make the comment
    try:
        raw = web.post('https://api.github.com/repos/' + gitAPI[1] + '/issues/'
                       + number + '/comments?access_token=' + gitAPI[0],
                       json.dumps({'body': '``\n' + post + '``'}))
    except OSError:  # HTTPError:
        return bot.say('The GitHub API returned an error.')

    data = json.loads(raw)
    bot.say('Added traceback to issue #%s. %s' % (number, data['html_url']))
    bot.debug(__file__, 'Traceback added to #%s in %s.' % (number, trigger.sender), 'warning')

@commands('findissue', 'findbug')
def findIssue(bot, trigger):
    """Search for a GitHub issue by keyword or ID. usage: .findissue search keywords/ID (optional) You can specify the first keyword as "CLOSED" to search closed issues."""
    if not trigger.group(2):
        return bot.reply('What are you searching for?')

    # Is the Oauth token and repo available?
    gitAPI = checkConfig(bot)
    if not gitAPI:
        return bot.say('Git module not configured, make sure github.oauth_token and github.repo are defined')
    firstParam = trigger.group(2).split(' ')[0]
    if firstParam.isdigit():
        URL = 'https://api.github.com/repos/%s/issues/%s' % (gitAPI[1], trigger.group(2))
    elif firstParam == 'CLOSED':
        if '%20'.join(trigger.group(2).split(' ')[1:]) not in ('', '\x02', '\x03'):
            URL = 'https://api.github.com/legacy/issues/search/' + gitAPI[1] + '/closed/' + '%20'.join(trigger.group(2).split(' ')[1:])
        else:
            return bot.reply('What are you searching for?')
    else:
        URL = 'https://api.github.com/legacy/issues/search/%s/open/%s' % (gitAPI[1], trigger.group(2))

    try:
        raw = web.get(URL)
    except HTTPError:
        return bot.say('The GitHub API returned an error.')

    try:
        if firstParam.isdigit():
            data = json.loads(raw)
        else:
            data = json.loads(raw)['issues'][-1]
    except (KeyError, IndexError):
        return bot.say('No search results.')
    try:
        if len(data['body'].split('\n')) > 1:
            body = data['body'].split('\n')[0] + '...'
        else:
            body = data['body'].split('\n')[0]
    except (KeyError):
        bot.debug(
            'GitHub KeyErr',
            ('API returned an invalid result on query request ' +
             trigger.group(2)),
            'always')
        return bot.say('Invalid result, please try again later.')
    bot.reply('[#%s]\x02title:\x02 %s \x02|\x02 %s' % (data['number'], data['title'], body))
    bot.say(data['html_url'])
<commit_msg>[github]Parse issue URLs as if they are .findissue [number]<commit_after>"""
github.py - Willie Github Module
Copyright 2012, Dimitri Molenaars http://tyrope.nl/
Licensed under the Eiffel Forum License 2.

http://willie.dftba.net/
"""

from datetime import datetime
from urllib2 import HTTPError
import json
from willie import web, tools
from willie.module import commands, rule, NOLIMIT
import os
import re

issueURL = r'https?://(?:www\.)?github.com/([A-z0-9\-]+/[A-z0-9\-]+)/issues/([\d]+)'

def checkConfig(bot):
    if not bot.config.has_option('github', 'oauth_token') or not bot.config.has_option('github', 'repo'):
        return False
    else:
        return [bot.config.github.oauth_token, bot.config.github.repo]

def configure(config):
    """
    | [github] | example | purpose |
    | -------- | ------- | ------- |
    | oauth_token | 5868e7af57496cc3ae255868e7af57496cc3ae25 | The OAuth token to connect to your github repo |
    | repo | embolalia/willie | The GitHub repo you're working from. |
    """
    chunk = ''
    if config.option('Configuring github issue reporting and searching module', False):
        config.interactive_add('github', 'oauth_token', 'Github API Oauth2 token', '')
        config.interactive_add('github', 'repo', 'Github repository', 'embolalia/willie')
    return chunk

def setup(bot):
    regex = re.compile(issueURL)
    if not bot.memory.contains('url_callbacks'):
        bot.memory['url_callbacks'] = tools.WillieMemory()
    bot.memory['url_callbacks'][regex] = issue_info

@commands('makeissue', 'makebug')
def issue(bot, trigger):
    """Create a GitHub issue, also known as a bug report. Syntax: .makeissue Title of the bug report"""
    # check input
    if not trigger.group(2):
        return bot.say('Please title the issue')

    # Is the Oauth token and repo available?
    gitAPI = checkConfig(bot)
    if not gitAPI:
        return bot.say('Git module not configured, make sure github.oauth_token and github.repo are defined')

    # parse input
    now = ' '.join(str(datetime.utcnow()).split(' ')).split('.')[0] + ' UTC'
    body = 'Submitted by: %s\nFrom channel: %s\nAt %s' % (trigger.nick, trigger.sender, now)
    data = {"title": trigger.group(2).encode('utf-8'), "body": body}
    # submit
    try:
        raw = web.post('https://api.github.com/repos/' + gitAPI[1] + '/issues?access_token=' + gitAPI[0], json.dumps(data))
    except HTTPError:
        bot.say('The GitHub API returned an error.')
        return NOLIMIT

    data = json.loads(raw)
    bot.say('Issue #%s posted. %s' % (data['number'], data['html_url']))
    bot.debug(__file__, 'Issue #%s created in %s' % (data['number'], trigger.sender), 'warning')

@commands('addtrace', 'addtraceback')
def add_traceback(bot, trigger):
    """Add a traceback to a GitHub issue.

    This pulls the traceback from the exceptions log file. To use, put .addtrace
    followed by the issue number to add the comment to, then the signature of
    the error (the message shown to the channel when the error occured). This
    command will only work for errors from unhandled exceptions."""
    # Make sure the API is set up
    gitAPI = checkConfig(bot)
    if not gitAPI:
        return bot.say('Git module not configured, make sure github.oauth_token and github.repo are defined')

    # Make sure the input is valid
    args = trigger.group(2).split(None, 1)
    if len(args) != 2:
        bot.say(validate)
        return
    number, trace = args

    # Make sure the given issue number exists
    issue_data = web.get('https://api.github.com/repos/%s/issues/%s' % (gitAPI[1], number))
    issue_data = json.loads(issue_data)
    if 'message' in issue_data and issue_data['message'] == 'Not Found':
        return bot.say("That issue doesn't exist.")

    # Find the relevant lines from the log file
    post = ''
    logfile = os.path.join(bot.config.logdir, 'exceptions.log')
    with open(logfile) as log:
        in_trace = False
        for data in log:
            if data == 'Signature: ' + trace + '\n':
                post = data
                in_trace = True
            elif data == '----------------------------------------\n':
                in_trace = False
            elif in_trace:
                post += data

    # Give an error if we didn't find the traceback
    if not post:
        return bot.say("I don't remember getting that error. Please post it "
                       "yourself at https://github.com/%s/issues/%s"
                       % (gitAPI[1], number))

    # Make the comment
    try:
        raw = web.post('https://api.github.com/repos/' + gitAPI[1] + '/issues/'
                       + number + '/comments?access_token=' + gitAPI[0],
                       json.dumps({'body': '``\n' + post + '``'}))
    except OSError:  # HTTPError:
        bot.say('The GitHub API returned an error.')
        return NOLIMIT

    data = json.loads(raw)
    bot.say('Added traceback to issue #%s. %s' % (number, data['html_url']))
    bot.debug(__file__, 'Traceback added to #%s in %s.' % (number, trigger.sender), 'warning')

@commands('findissue', 'findbug')
def findIssue(bot, trigger):
    """Search for a GitHub issue by keyword or ID. usage: .findissue search keywords/ID (optional) You can specify the first keyword as "CLOSED" to search closed issues."""
    if not trigger.group(2):
        return bot.reply('What are you searching for?')

    # Is the Oauth token and repo available?
    gitAPI = checkConfig(bot)
    if not gitAPI:
        return bot.say('Git module not configured, make sure github.oauth_token and github.repo are defined')
    firstParam = trigger.group(2).split(' ')[0]
    if firstParam.isdigit():
        URL = 'https://api.github.com/repos/%s/issues/%s' % (gitAPI[1], trigger.group(2))
    elif firstParam == 'CLOSED':
        if '%20'.join(trigger.group(2).split(' ')[1:]) not in ('', '\x02', '\x03'):
            URL = 'https://api.github.com/legacy/issues/search/' + gitAPI[1] + '/closed/' + '%20'.join(trigger.group(2).split(' ')[1:])
        else:
            return bot.reply('What are you searching for?')
    else:
        URL = 'https://api.github.com/legacy/issues/search/%s/open/%s' % (gitAPI[1], trigger.group(2))

    try:
        raw = web.get(URL)
    except HTTPError:
        bot.say('The GitHub API returned an error.')
        return NOLIMIT

    try:
        if firstParam.isdigit():
            data = json.loads(raw)
        else:
            data = json.loads(raw)['issues'][-1]
    except (KeyError, IndexError):
        return bot.say('No search results.')
    try:
        if len(data['body'].split('\n')) > 1:
            body = data['body'].split('\n')[0] + '...'
        else:
            body = data['body'].split('\n')[0]
    except (KeyError):
        bot.debug(
            'GitHub KeyErr',
            ('API returned an invalid result on query request ' +
             trigger.group(2)),
            'always')
        bot.say('Invalid result, please try again later.')
        return NOLIMIT
    bot.reply('[#%s]\x02title:\x02 %s \x02|\x02 %s' % (data['number'], data['title'], body))
    bot.say(data['html_url'])

@rule('.*%s.*' % issueURL)
def issue_info(bot, trigger, match=None):
    match = match or trigger
    URL = 'https://api.github.com/repos/%s/issues/%s' % (match.group(1), match.group(2))

    try:
        raw = web.get(URL)
    except HTTPError:
        bot.say('The GitHub API returned an error.')
        return NOLIMIT
    data = json.loads(raw)
    try:
        if len(data['body'].split('\n')) > 1:
            body = data['body'].split('\n')[0] + '...'
        else:
            body = data['body'].split('\n')[0]
    except (KeyError):
        bot.say('The API says this is an invalid issue. Please report this if you know it\'s a correct link!')
        return NOLIMIT
    bot.reply('[#%s]\x02title:\x02 %s \x02|\x02 %s' % (data['number'], data['title'], body))

<|endoftext|>


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

from optparse import OptionParser

import json
import sys
import urllib2

def main():
  parser = OptionParser(usage="check_spark [spark_master_hostname] [num_workers]",
      add_help_option=True)
  (opts, args) = parser.parse_args()
  if len(args) != 2:
    parser.print_help()
    sys.exit(1)
  else:
    status = check_spark_master(args[0], args[1])
    sys.exit(status)

def check_spark_master(master_hostname, num_workers):
  master_url = "http://" + master_hostname + ":8080"
  url = master_url + "?format=json"
  response = urllib2.urlopen(url, timeout=30)
  if response.code != 200:
    print("Spark master " + url + " returned " + str(response.code))
    return 1
  master_json = response.read()
  return check_spark_json(master_json, num_workers)

def check_spark_json(spark_json, num_workers):
  json_data = json.loads(spark_json)
  ## Find number of workers from status page
  got_workers = len(json_data.get("workers"))
  if got_workers == int(num_workers):
    return 0
  else:
    return 1

if __name__ == "__main__":
  main()
<commit_msg>Remove unused spark-standalone health check script<commit_after><|endoftext|>


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

import socket
import sys
import argparse
from datetime import datetime
from threading import Thread

HOST = ''		# Symbolic name meaning all available interfaces
PORT = 6666		# Default netconsole client IN Port

# Netconsole Client
class Netconsole():
	def __init__(self, host, port):
		self.host = host
		self.port = port
		self.max_size = 4096
		self.client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
		self.client_addr = (host, port)
		self.client.bind(self.client_addr)
		self.stopflag = 0
		self.server = None

	def _recv(self):
		while True:
			self.data, self.server = self.client.recvfrom(self.max_size)
			if args.HOST != '':
				if self.server[0] == args.HOST:
					sys.stdout.write(self.data)
			else:
				sys.stdout.write(self.data)

	def _send(self):
		while True:
			self.client_input = sys.stdin.readline()
			if self.client_input == "quit\n":
				if self.server is not None:
					self.client.sendto("Leave Netconsole Client.\n", self.server)
				break
			if self.client_input is not None:
				if self.server is not None:
					self.client.sendto(self.client_input, self.server)

	def worker_send(self):
		th_send = Thread(target = self._send)
		th_send.setDaemon(True)
		th_send.start()
		th_send.join()
		self.stopflag = 1

	def worker_recv(self):
		th_recv = Thread(target = self._recv)
		th_recv.setDaemon(True)
		th_recv.start()

	def close(self):
		while self.stopflag == 0 :
			pass
		self.client.close()

def argv_gen():
	parser = argparse.ArgumentParser(description='A Python Based Netconsole Client')
	parser.add_argument('--ip', type=str, action="store", dest='HOST',
			default='', help="Specify ip address to listen. Default is ''")
	parser.add_argument('-p', '--port', type=int, action="store", dest='PORT',
			default=6666, help="Specify port to listen. Default 6666 is used.")
	args = parser.parse_args()
	return args

if __name__ == "__main__" :
	args = argv_gen()
	client1 = Netconsole(HOST, args.PORT)
	if args.HOST == '':
		print "Listen to port %d" % args.PORT
	else :
		print "Listen to port %d, target is %s" % (args.PORT, args.HOST)
	client1.worker_recv()
	client1.worker_send()
	client1.close()
<commit_msg>Marked send() and recv() as private methods<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import socket
import sys
import argparse
from datetime import datetime
from threading import Thread

HOST = ''		# Symbolic name meaning all available interfaces
PORT = 6666		# Default netconsole client IN Port

# Netconsole Client
class Netconsole():
	def __init__(self, host, port):
		self.host = host
		self.port = port
		self.max_size = 4096
		self.client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
		self.client_addr = (host, port)
		self.client.bind(self.client_addr)
		self.stopflag = 0
		self.server = None

	def __recv(self):
		while True:
			self.data, self.server = self.client.recvfrom(self.max_size)
			if args.HOST != '':
				if self.server[0] == args.HOST:
					sys.stdout.write(self.data)
			else:
				sys.stdout.write(self.data)

	def __send(self):
		while True:
			self.client_input = sys.stdin.readline()
			if self.client_input == "quit\n":
				if self.server is not None:
					self.client.sendto("Leave Netconsole Client.\n", self.server)
				break
			if self.client_input is not None:
				if self.server is not None:
					self.client.sendto(self.client_input, self.server)

	def worker_send(self):
		th_send = Thread(target = self.__send)
		th_send.setDaemon(True)
		th_send.start()
		th_send.join()
		self.stopflag = 1

	def worker_recv(self):
		th_recv = Thread(target = self.__recv)
		th_recv.setDaemon(True)
		th_recv.start()

	def close(self):
		while self.stopflag == 0 :
			pass
		self.client.close()

def argv_gen():
	parser = argparse.ArgumentParser(description='A Python Based Netconsole Client')
	parser.add_argument('--ip', type=str, action="store", dest='HOST',
			default='', help="Specify ip address to listen. Default is ''")
	parser.add_argument('-p', '--port', type=int, action="store", dest='PORT',
			default=6666, help="Specify port to listen. Default 6666 is used.")
	args = parser.parse_args()
	return args

if __name__ == "__main__" :
	args = argv_gen()
	client1 = Netconsole(HOST, args.PORT)
	if args.HOST == '':
		print "Listen to port %d" % args.PORT
	else :
		print "Listen to port %d, target is %s" % (args.PORT, args.HOST)
	client1.worker_recv()
	client1.worker_send()
	client1.close()
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2013 The Plaso Project Authors.
# Please see the AUTHORS file for details on individual authors.
#
# 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.
"""A simple script to export files out of an image."""

import argparse
import hashlib
import logging
import os
import sys

from dfvfs.lib import definitions
from dfvfs.path import factory as path_spec_factory
from dfvfs.resolver import resolver as path_spec_resolver

from plaso import preprocessors
from plaso.collector import collector
from plaso.frontend import utils as frontend_utils
from plaso.lib import errors
from plaso.lib import event
from plaso.lib import preprocess_interface
from plaso.lib import queue

class ImageExtractor(object):
  """Class that implements the image extractor."""

  def __init__(self):
    """Initializes the image extractor object."""
    super(ImageExtractor, self).__init__()
    self._image_path = None
    self._image_offset = 0
    self._image_collector = None
    self._pre_obj = None

  def Open(self, image_path, sector_offset=0):
    """Opens the image.

    Args:
      image_path: The path to the disk image to parse.
      sector_offset: Option offset in sectors into the image where the partition
                     containing the file system starts. The default is 0.
    """
    self._image_path = image_path
    self._image_offset = sector_offset * 512

  def _Preprocess(self):
    """Preprocesses the image.

    Returns:
      The image collector object.
    """
    if self._pre_obj is not None:
      return

    self._pre_obj = event.PreprocessObject()
    try:
      image_collector = collector.GenericPreprocessCollector(
          self._pre_obj, self._image_path)
      image_collector.SetImageInformation(byte_offset=self._image_offset)

    except errors.UnableToOpenFilesystem as e:
      raise RuntimeError(
          u'Unable to proceed, not an image file? [{0:s}]'.format(e))

    plugin_list = preprocessors.PreProcessList(
        self._pre_obj, image_collector)

    logging.info(u'Guessing OS')
    guessed_os = preprocess_interface.GuessOS(image_collector)
    logging.info(u'OS: {0:s}'.format(guessed_os))

    logging.info(u'Running preprocess.')
    for weight in plugin_list.GetWeightList(guessed_os):
      for plugin in plugin_list.GetWeight(guessed_os, weight):
        try:
          plugin.Run()
        except errors.PreProcessFail as e:
          logging.warning(
              u'Unable to run preprocessor: {0:s} with error: {1:s}'.format(
                  plugin.__class__.__name__, e))

    logging.info(u'Preprocess done, saving files from image.')

    return image_collector

  def ExtractWithFilter(
      self, filter_file_path, destination_path, process_vss=False,
      remove_duplicates=False):
    """Extracts files using a filter expression.

    This method runs the file extraction process on the image and
    potentially on every VSS if that is wanted.

    Args:
      filter_file_path: The path of the file that contains the filter
                        expressions.
      destination_path: The path where the extracted files should be stored.
      process_vss: Optional boolean value to indicate if VSS should be detected
                   and processed. The default is false.
      remove_duplicates: Optional boolean value to indicate if duplicates should
                         be detected and removed. The default is false.
    """
    if self._pre_obj is None:
      image_collector = self._Preprocess()

    if not os.path.isdir(destination_path):
      os.makedirs(destination_path)

    # Save the regular files.
    FileSaver.calc_md5 = remove_duplicates
    filter_object = collector.BuildCollectionFilterFromFile(filter_file_path)

    for path_spec in image_collector.GetPathSpecs(filter_object):
      FileSaver.WriteFile(path_spec, destination_path)

    if process_vss:
      logging.info(u'Extracting files from VSS.')
      os_path_spec = path_spec_factory.Factory.NewPathSpec(
          definitions.TYPE_INDICATOR_OS, location=self._image_path)

      if self._image_offset > 0:
        volume_path_spec = path_spec_factory.Factory.NewPathSpec(
          definitions.TYPE_INDICATOR_TSK_PARTITION,
          start_offset=self._image_offset, parent=os_path_spec)
      else:
        volume_path_spec = os_path_spec

      vss_path_spec = path_spec_factory.Factory.NewPathSpec(
          definitions.TYPE_INDICATOR_VSHADOW, location=u'/',
          parent=volume_path_spec)

      vss_file_entry = path_spec_resolver.Resolver.OpenFileEntry(vss_path_spec)

      number_of_vss = vss_file_entry.number_of_sub_file_entries

      for store_number in range(0, number_of_vss):
        logging.info(u'Extracting files from VSS {0:d} out of {1:d}'.format(
            store_number + 1, number_of_vss))

        vss_collector = collector.GenericPreprocessCollector(
            self._pre_obj, self._image_path)
        vss_collector.SetImageInformation(byte_offset=self._image_offset)
        vss_collector.SetVssInformation(store_index=store_number)

        filename_prefix = 'vss_{0:d}'.format(store_number)
        filter_object = collector.BuildCollectionFilterFromFile(
            filter_file_path)

        for path_spec in vss_collector.GetPathSpecs(filter_object):
          FileSaver.WriteFile(
              path_spec, destination_path, filename_prefix=filename_prefix)

  def ExtractWithExtensions(
       self, extensions, destination_path, process_vss=False,
       remove_duplicates=False):
    """Extracts files using extensions.

    Args:
      extensions: a list of extensions.
      destination_path: the path where the extracted files should be stored.
      process_vss: optional boolean value to indicate if VSS should be detected
                   and processed. The default is false.
      remove_duplicates: optional boolean value to indicate if duplicates should
                         be detected and removed. The default is false.
    """
    if not os.path.isdir(destination_path):
      os.makedirs(destination_path)

    input_queue = queue.SingleThreadedQueue()
    output_queue = queue.SingleThreadedQueue()

    image_collector = collector.Collector(
        input_queue, output_queue, self._image_path)
    image_collector.SetImageInformation(byte_offset=self._image_offset)
    if process_vss:
      # TODO: don't we need a store number here?
      image_collector.SetVssInformation()

    image_collector.Collect()

    FileSaver.calc_md5 = remove_duplicates

    input_queue_consumer = ImageExtractorQueueConsumer(
        input_queue, extensions, destination_path)
    input_queue_consumer.ConsumePathSpecs()

class ImageExtractorQueueConsumer(queue.PathSpecQueueConsumer):
  """Class that implements an image extractor queue consumer."""

  def __init__(self, process_queue, extensions, destination_path):
    """Initializes the image extractor queue consumer.

    Args:
      process_queue: the process queue (instance of Queue).
      extensions: a list of extensions.
      destination_path: the path where the extracted files should be stored.
    """
    super(ImageExtractorQueueConsumer, self).__init__(process_queue)
    self._destination_path = destination_path
    self._extensions = extensions

  def _ConsumePathSpec(self, path_spec):
    """Consumes a path specification callback for ConsumePathSpecs."""
    # TODO: move this into a function of path spec e.g. GetExtension().
    _, _, extension = path_spec.file_path.rpartition('.')
    if extension.lower() in self._extensions:
      vss_store_number = getattr(path_spec, 'vss_store_number', None)
      if vss_store_number is not None:
        filename_prefix = 'vss_{0:d}'.format(vss_store_number + 1)
      else:
        filename_prefix = ''

      FileSaver.WriteFile(
          path_spec, self._destination_path, filename_prefix=filename_prefix)

class FileSaver(object):
  """A simple class that is used to save files."""
  md5_dict = {}
  calc_md5 = False

  @classmethod
  def WriteFile(cls, source_path_spec, destination_path, filename_prefix=''):
    """Writes the contents of the source to the destination file.

    Args:
      source_path_spec: the path specification of the source file.
      destination_path: the path of the destination file.
      filename_prefix: optional prefix for the filename. The default is an
                       empty string.
    """
    file_entry = path_spec_resolver.Resolver.OpenFileEntry(source_path_spec)
    directory = u''
    filename = file_entry.name

    # There will be issues on systems that use a different separator than a
    # forward slash. However a forward slash is always used in the pathspec.
    if os.path.sep != u'/':
      filename = filename.replace(u'/', os.path.sep)

    if os.path.sep in filename:
      directory_string, _, filename = filename.rpartition(os.path.sep)
      if directory_string:
        directory = os.path.join(
            destination_path, *directory_string.split(os.path.sep))

    if filename_prefix:
      extracted_filename = u'{0:s}_{1:s}'.format(filename_prefix, filename)
    else:
      extracted_filename = filename

    while extracted_filename.startswith(os.path.sep):
      extracted_filename = extracted_filename[1:]

    if directory:
      if not os.path.isdir(directory):
        os.makedirs(directory)
    else:
      directory = destination_path

    save_file = True
    if cls.calc_md5:
      stat = file_entry.GetStat()
      inode = stat.attributes.get('ino', 0)
      md5sum = CalculateHash(file_entry)
      if inode in cls.md5_dict:
        if md5sum in cls.md5_dict[inode]:
          save_file = False
        else:
          cls.md5_dict[inode].append(md5sum)
      else:
        cls.md5_dict[inode] = [md5sum]

    if save_file:
      try:
        frontend_utils.OutputWriter.WriteFile(
            file_entry, os.path.join(directory, extracted_filename))
      except IOError as exception:
        logging.error(
            u'[skipping] unable to save file: {0:s} with error: {1:s}'.format(
                filename, exception))

def CalculateHash(file_object):
  """Return a hash for a given file object."""
  md5 = hashlib.md5()
  file_object.seek(0)

  data = file_object.read(4098)
  while data:
    md5.update(data)
    data = file_object.read(4098)

  return md5.hexdigest()

def Main():
  """The main function, running the show."""
  arg_parser = argparse.ArgumentParser(
      description=(
          'This is a simple collector designed to export files inside an '
          'image, both within a regular RAW image as well as inside a VSS. '
          'The tool uses a collection filter that uses the same syntax as a '
          'targeted plaso filter.'),
      epilog='And that\'s how you export files, plaso style.')

  arg_parser.add_argument(
      '-v', '--vss', dest='vss', action='store_true', default=False,
      help='Also extract files from VSS.')

  arg_parser.add_argument(
      '-o', '--offset', dest='offset', action='store', default=0, type=int,
      help='Offset in sectors into where the partition starts.')

  arg_parser.add_argument(
      '-w', '--write', dest='path', action='store', default='.', type=str,
      help='The directory in which extracted files should be stored in.')

  arg_parser.add_argument(
      '-x', '--extensions', dest='extension_string', action='store',
      type=str, metavar='EXTENSION_STRING', help=(
          'If the purpose is to find all files given a certain extension '
          'this options should be used. This option accepts a comma separated '
          'string denoting all file extensions, eg: -x "csv,docx,pst".'))

  arg_parser.add_argument(
      '-f', '--filter', action='store', dest='filter', metavar='FILTER_FILE',
      type=str, help=(
          'Full path to the file that contains the collection filter, '
          'the file can use variables that are defined in preprocesing, '
          'just like any other log2timeline/plaso collection filter.'))

  arg_parser.add_argument(
      '--include_duplicates', dest='include_duplicates', action='store_true',
      default=False, help=(
          'By default if VSS is turned on all files saved will have their '
          'MD5 sum calculated and compared to other files already saved '
          'with the same inode value. If the MD5 sum is the same the file '
          'does not get saved again. This option turns off that behavior '
          'so that all files will get stored, even if they are duplicates.'))

  arg_parser.add_argument(
      'image', action='store', metavar='IMAGE', default=None, type=str, help=(
          'The full path to the image file that we are about to extract files '
          'from, it should be a raw image or another image that plaso '
          'supports.'))

  options = arg_parser.parse_args()

  format_str = '[%(levelname)s] %(message)s'
  logging.basicConfig(level=logging.INFO, format=format_str)

  if not os.path.isfile(options.image):
    raise RuntimeError(u'Unable to proceed, image file does not exist.')

  if not (options.filter or options.extension_string):
    arg_parser.print_help()
    print ''
    logging.error(u'Neither extension string nor filter defined.')
    sys.exit(1)

  if options.filter:
    if not os.path.isfile(options.filter):
      arg_parser.print_help()
      print ''
      logging.error(u'Unable to proceed, filter file does not exist.')
      sys.exit(1)

  if not options.vss:
    options.remove_duplicates = False
  elif options.include_duplicates:
    options.remove_duplicates = False
  else:
    options.remove_duplicates = True

  image_extractor = ImageExtractor()
  image_extractor.Open(options.image, sector_offset=options.offset)

  # The option of running both options.
  if options.filter:
    image_extractor.ExtractWithFilter(
        options.filter, options.path,
        process_vss=options.vss,
        remove_duplicates=options.remove_duplicates)

  if options.extension_string:
    extensions = [x.strip() for x in options.extension_string.split(',')]

    image_extractor.ExtractWithExtensions(
        extensions, options.path, process_vss=options.vss,
        remove_duplicates=options.remove_duplicates)
    logging.info(u'Files based on extension extracted.')

if __name__ == '__main__':
  Main()
<commit_msg>Revert "Fix parameters argument inversion"<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2013 The Plaso Project Authors.
# Please see the AUTHORS file for details on individual authors.
#
# 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.
"""A simple script to export files out of an image."""

import argparse
import hashlib
import logging
import os
import sys

from dfvfs.lib import definitions
from dfvfs.path import factory as path_spec_factory
from dfvfs.resolver import resolver as path_spec_resolver

from plaso import preprocessors
from plaso.collector import collector
from plaso.frontend import utils as frontend_utils
from plaso.lib import errors
from plaso.lib import event
from plaso.lib import preprocess_interface
from plaso.lib import queue

class ImageExtractor(object):
  """Class that implements the image extractor."""

  def __init__(self):
    """Initializes the image extractor object."""
    super(ImageExtractor, self).__init__()
    self._image_path = None
    self._image_offset = 0
    self._image_collector = None
    self._pre_obj = None

  def Open(self, image_path, sector_offset=0):
    """Opens the image.

    Args:
      image_path: The path to the disk image to parse.
      sector_offset: Option offset in sectors into the image where the partition
                     containing the file system starts. The default is 0.
    """
    self._image_path = image_path
    self._image_offset = sector_offset * 512

  def _Preprocess(self):
    """Preprocesses the image.

    Returns:
      The image collector object.
    """
    if self._pre_obj is not None:
      return

    self._pre_obj = event.PreprocessObject()
    try:
      image_collector = collector.GenericPreprocessCollector(
          self._pre_obj, self._image_path)
      image_collector.SetImageInformation(byte_offset=self._image_offset)

    except errors.UnableToOpenFilesystem as e:
      raise RuntimeError(
          u'Unable to proceed, not an image file? [{0:s}]'.format(e))

    plugin_list = preprocessors.PreProcessList(
        self._pre_obj, image_collector)

    logging.info(u'Guessing OS')
    guessed_os = preprocess_interface.GuessOS(image_collector)
    logging.info(u'OS: {0:s}'.format(guessed_os))

    logging.info(u'Running preprocess.')
    for weight in plugin_list.GetWeightList(guessed_os):
      for plugin in plugin_list.GetWeight(guessed_os, weight):
        try:
          plugin.Run()
        except errors.PreProcessFail as e:
          logging.warning(
              u'Unable to run preprocessor: {0:s} with error: {1:s}'.format(
                  plugin.__class__.__name__, e))

    logging.info(u'Preprocess done, saving files from image.')

    return image_collector

  def ExtractWithFilter(
      self, filter_file_path, destination_path, process_vss=False,
      remove_duplicates=False):
    """Extracts files using a filter expression.

    This method runs the file extraction process on the image and
    potentially on every VSS if that is wanted.

    Args:
      filter_file_path: The path of the file that contains the filter
                        expressions.
      destination_path: The path where the extracted files should be stored.
      process_vss: Optional boolean value to indicate if VSS should be detected
                   and processed. The default is false.
      remove_duplicates: Optional boolean value to indicate if duplicates should
                         be detected and removed. The default is false.
    """
    if self._pre_obj is None:
      image_collector = self._Preprocess()

    if not os.path.isdir(destination_path):
      os.makedirs(destination_path)

    # Save the regular files.
    FileSaver.calc_md5 = remove_duplicates
    filter_object = collector.BuildCollectionFilterFromFile(filter_file_path)

    for path_spec in image_collector.GetPathSpecs(filter_object):
      FileSaver.WriteFile(path_spec, destination_path)

    if process_vss:
      logging.info(u'Extracting files from VSS.')
      os_path_spec = path_spec_factory.Factory.NewPathSpec(
          definitions.TYPE_INDICATOR_OS, location=self._image_path)

      if self._image_offset > 0:
        volume_path_spec = path_spec_factory.Factory.NewPathSpec(
          definitions.TYPE_INDICATOR_TSK_PARTITION,
          start_offset=self._image_offset, parent=os_path_spec)
      else:
        volume_path_spec = os_path_spec

      vss_path_spec = path_spec_factory.Factory.NewPathSpec(
          definitions.TYPE_INDICATOR_VSHADOW, location=u'/',
          parent=volume_path_spec)

      vss_file_entry = path_spec_resolver.Resolver.OpenFileEntry(vss_path_spec)

      number_of_vss = vss_file_entry.number_of_sub_file_entries

      for store_number in range(0, number_of_vss):
        logging.info(u'Extracting files from VSS {0:d} out of {1:d}'.format(
            store_number + 1, number_of_vss))

        vss_collector = collector.GenericPreprocessCollector(
            self._pre_obj, self._image_path)
        vss_collector.SetImageInformation(byte_offset=self._image_offset)
        vss_collector.SetVssInformation(store_index=store_number)

        filename_prefix = 'vss_{0:d}'.format(store_number)
        filter_object = collector.BuildCollectionFilterFromFile(
            filter_file_path)

        for path_spec in vss_collector.GetPathSpecs(filter_object):
          FileSaver.WriteFile(
              path_spec, destination_path, filename_prefix=filename_prefix)

  def ExtractWithExtensions(
       self, extensions, destination_path, process_vss=False,
       remove_duplicates=False):
    """Extracts files using extensions.

    Args:
      extensions: a list of extensions.
      destination_path: the path where the extracted files should be stored.
      process_vss: optional boolean value to indicate if VSS should be detected
                   and processed. The default is false.
      remove_duplicates: optional boolean value to indicate if duplicates should
                         be detected and removed. The default is false.
    """
    if not os.path.isdir(destination_path):
      os.makedirs(destination_path)

    input_queue = queue.SingleThreadedQueue()
    output_queue = queue.SingleThreadedQueue()

    image_collector = collector.Collector(
        input_queue, output_queue, self._image_path)
    image_collector.SetImageInformation(byte_offset=self._image_offset)
    if process_vss:
      # TODO: don't we need a store number here?
      image_collector.SetVssInformation()

    image_collector.Collect()

    FileSaver.calc_md5 = remove_duplicates

    input_queue_consumer = ImageExtractorQueueConsumer(
        input_queue, extensions, destination_path)
    input_queue_consumer.ConsumePathSpecs()

class ImageExtractorQueueConsumer(queue.PathSpecQueueConsumer):
  """Class that implements an image extractor queue consumer."""

  def __init__(self, process_queue, extensions, destination_path):
    """Initializes the image extractor queue consumer.

    Args:
      process_queue: the process queue (instance of Queue).
      extensions: a list of extensions.
      destination_path: the path where the extracted files should be stored.
    """
    super(ImageExtractorQueueConsumer, self).__init__(process_queue)
    self._destination_path = destination_path
    self._extensions = extensions

  def _ConsumePathSpec(self, path_spec):
    """Consumes a path specification callback for ConsumePathSpecs."""
    # TODO: move this into a function of path spec e.g. GetExtension().
    _, _, extension = path_spec.file_path.rpartition('.')
    if extension.lower() in self._extensions:
      vss_store_number = getattr(path_spec, 'vss_store_number', None)
      if vss_store_number is not None:
        filename_prefix = 'vss_{0:d}'.format(vss_store_number + 1)
      else:
        filename_prefix = ''

      FileSaver.WriteFile(
          path_spec, self._destination_path, filename_prefix=filename_prefix)

class FileSaver(object):
  """A simple class that is used to save files."""
  md5_dict = {}
  calc_md5 = False

  @classmethod
  def WriteFile(cls, source_path_spec, destination_path, filename_prefix=''):
    """Writes the contents of the source to the destination file.

    Args:
      source_path_spec: the path specification of the source file.
      destination_path: the path of the destination file.
      filename_prefix: optional prefix for the filename. The default is an
                       empty string.
    """
    file_entry = path_spec_resolver.Resolver.OpenFileEntry(source_path_spec)
    directory = u''
    filename = file_entry.name

    # There will be issues on systems that use a different separator than a
    # forward slash. However a forward slash is always used in the pathspec.
    if os.path.sep != u'/':
      filename = filename.replace(u'/', os.path.sep)

    if os.path.sep in filename:
      directory_string, _, filename = filename.rpartition(os.path.sep)
      if directory_string:
        directory = os.path.join(
            destination_path, *directory_string.split(os.path.sep))

    if filename_prefix:
      extracted_filename = u'{0:s}_{1:s}'.format(filename_prefix, filename)
    else:
      extracted_filename = filename

    while extracted_filename.startswith(os.path.sep):
      extracted_filename = extracted_filename[1:]

    if directory:
      if not os.path.isdir(directory):
        os.makedirs(directory)
    else:
      directory = destination_path

    save_file = True
    if cls.calc_md5:
      stat = file_entry.GetStat()
      inode = stat.attributes.get('ino', 0)
      md5sum = CalculateHash(file_entry)
      if inode in cls.md5_dict:
        if md5sum in cls.md5_dict[inode]:
          save_file = False
        else:
          cls.md5_dict[inode].append(md5sum)
      else:
        cls.md5_dict[inode] = [md5sum]

    if save_file:
      try:
        frontend_utils.OutputWriter.WriteFile(
            file_entry, os.path.join(directory, extracted_filename))
      except IOError as exception:
        logging.error(
            u'[skipping] unable to save file: {0:s} with error: {1:s}'.format(
                filename, exception))

def CalculateHash(file_object):
  """Return a hash for a given file object."""
  md5 = hashlib.md5()
  file_object.seek(0)

  data = file_object.read(4098)
  while data:
    md5.update(data)
    data = file_object.read(4098)

  return md5.hexdigest()

def Main():
  """The main function, running the show."""
  arg_parser = argparse.ArgumentParser(
      description=(
          'This is a simple collector designed to export files inside an '
          'image, both within a regular RAW image as well as inside a VSS. '
          'The tool uses a collection filter that uses the same syntax as a '
          'targeted plaso filter.'),
      epilog='And that\'s how you export files, plaso style.')

  arg_parser.add_argument(
      '-v', '--vss', dest='vss', action='store_true', default=False,
      help='Also extract files from VSS.')

  arg_parser.add_argument(
      '-o', '--offset', dest='offset', action='store', default=0, type=int,
      help='Offset in sectors into where the partition starts.')

  arg_parser.add_argument(
      '-w', '--write', dest='path', action='store', default='.', type=str,
      help='The directory in which extracted files should be stored in.')

  arg_parser.add_argument(
      '-x', '--extensions', dest='extension_string', action='store',
      type=str, metavar='EXTENSION_STRING', help=(
          'If the purpose is to find all files given a certain extension '
          'this options should be used. This option accepts a comma separated '
          'string denoting all file extensions, eg: -x "csv,docx,pst".'))

  arg_parser.add_argument(
      '-f', '--filter', action='store', dest='filter', metavar='FILTER_FILE',
      type=str, help=(
          'Full path to the file that contains the collection filter, '
          'the file can use variables that are defined in preprocesing, '
          'just like any other log2timeline/plaso collection filter.'))

  arg_parser.add_argument(
      '--include_duplicates', dest='include_duplicates', action='store_true',
      default=False, help=(
          'By default if VSS is turned on all files saved will have their '
          'MD5 sum calculated and compared to other files already saved '
          'with the same inode value. If the MD5 sum is the same the file '
          'does not get saved again. This option turns off that behavior '
          'so that all files will get stored, even if they are duplicates.'))

  arg_parser.add_argument(
      'image', action='store', metavar='IMAGE', default=None, type=str, help=(
          'The full path to the image file that we are about to extract files '
          'from, it should be a raw image or another image that plaso '
          'supports.'))

  options = arg_parser.parse_args()

  format_str = '[%(levelname)s] %(message)s'
  logging.basicConfig(level=logging.INFO, format=format_str)

  if not os.path.isfile(options.image):
    raise RuntimeError(u'Unable to proceed, image file does not exist.')

  if not (options.filter or options.extension_string):
    arg_parser.print_help()
    print ''
    logging.error(u'Neither extension string nor filter defined.')
    sys.exit(1)

  if options.filter:
    if not os.path.isfile(options.filter):
      arg_parser.print_help()
      print ''
      logging.error(u'Unable to proceed, filter file does not exist.')
      sys.exit(1)

  if not options.vss:
    options.remove_duplicates = False
  elif options.include_duplicates:
    options.remove_duplicates = False
  else:
    options.remove_duplicates = True

  image_extractor = ImageExtractor()
  image_extractor.Open(options.image, sector_offset=options.offset)

  # The option of running both options.
  if options.filter:
    image_extractor.ExtractWithFilter(
        options.path, options.filter,
        process_vss=options.vss,
        remove_duplicates=options.remove_duplicates)

  if options.extension_string:
    extensions = [x.strip() for x in options.extension_string.split(',')]

    image_extractor.ExtractWithExtensions(
        extensions, options.path, process_vss=options.vss,
        remove_duplicates=options.remove_duplicates)
    logging.info(u'Files based on extension extracted.')

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


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
# Copyright 2013 The Plaso Project Authors.
# Please see the AUTHORS file for details on individual authors.
#
# 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.
"""A simple script to export files out of an image."""
import argparse
import hashlib
import os
import logging
import sys

from plaso import preprocessors

from plaso.collector import factory as collector_factory
from plaso.lib import collector_filter
from plaso.lib import errors
from plaso.lib import event
from plaso.lib import pfile
from plaso.lib import preprocess
from plaso.lib import putils
from plaso.lib import queue
from plaso.lib import vss

def RunExtensionExtraction(options, extensions, fscache):
  """Extracts files based on file extensions.

  Args:
    options: A configuration object, an argparse like object.
    fscache: A FileSystemCache object.
  """
  input_queue = queue.SingleThreadedQueue()
  output_queue = queue.SingleThreadedQueue()

  my_collector = collector_factory.GetImageCollector(
      input_queue, output_queue, options.image, sector_offset=options.offset,
      parse_vss=options.vss, fscache=fscache)

  my_collector.Collect()

  # Configure filesaver.
  FileSaver.calc_md5 = options.remove_duplicates

  for pathspec_string in input_queue.PopItems():
    pathspec = event.EventPathSpec()
    pathspec.FromProtoString(pathspec_string)

    _, _, extension = pathspec.file_path.rpartition('.')
    if extension.lower() in extensions:
      fh = pfile.OpenPFile(pathspec, fscache=fscache)
      if getattr(pathspec, 'vss_store_number', None):
        FileSaver.prefix = 'vss_%d' % pathspec.vss_store_number + 1
      else:
        FileSaver.prefix = ''
      if os.path.sep != '/':
        fh.name = fh.name.replace('/', os.path.sep)

      FileSaver.SaveFile(fh, options.path)

def ExtractFiles(
    my_collector, filter_definition, export_path, fscache=None):
  """Extracts files that match the filter.

  Args:
    my_collector: A collection object that can be used to locate files.
    filter_definition: This can be either a PathFilter protobuf,
    a path to a file with ASCII presentation of the protobuf or a list
    of entries, where each entry is a single path.
    export_path: The path to where the files should be stored.
    fscache: A filesystem cache object, if applicable.

  Returns:
    None, unless calc_md5 is set then it returns a dict with MD5 sums.
  """
  col_fil = collector_filter.CollectionFilter(my_collector, filter_definition)

  for pathspec_string in col_fil.GetPathSpecs():
    with pfile.OpenPFile(pathspec_string, fscache=fscache) as fh:
      # There will be issues on systems that use a different separator than a
      # forward slash. However a forward slash is always used in the pathspec.
      if os.path.sep != '/':
        fh.name = fh.name.replace('/', os.path.sep)
      FileSaver.SaveFile(fh, export_path)

class FileSaver(object):
  """A simple class that is used to save files."""
  md5_dict = {}
  calc_md5 = False
  prefix = ''

  @classmethod
  def SaveFile(cls, fh, export_path):
    """Take a filehandle and an export path and save the file."""
    directory = ''
    filename = ''
    if os.path.sep in fh.name:
      directory_string, _, filename = fh.name.rpartition(os.path.sep)
      if directory_string:
        directory = os.path.join(
            export_path, *directory_string.split(os.path.sep))
    else:
      filename = fh.name

    if cls.prefix:
      extracted_filename = '{}_{}'.format(cls.prefix, filename)
    else:
      extracted_filename = filename

    while extracted_filename.startswith(os.path.sep):
      extracted_filename = extracted_filename[1:]

    if directory:
      if not os.path.isdir(directory):
        os.makedirs(directory)
    else:
      directory = export_path

    save_file = True
    if cls.calc_md5:
      stat = fh.Stat()
      inode = stat.attributes.get('ino', 0)
      md5sum = CalculateHash(fh)
      if inode in cls.md5_dict:
        if md5sum in cls.md5_dict[inode]:
          save_file = False
        else:
          cls.md5_dict[inode].append(md5sum)
      else:
        cls.md5_dict[inode] = [md5sum]

    if save_file:
      try:
        putils.Pfile2File(fh, os.path.join(directory, extracted_filename))
      except IOError as e:
        logging.error(u'Unable to save file: {} [skipping]. Error {}'.format(
            fh.name, e))

def CalculateHash(file_object):
  """Return a hash for a given file object."""
  md5 = hashlib.md5()
  file_object.seek(0)

  data = file_object.read(4098)
  while data:
    md5.update(data)
    data = file_object.read(4098)

  return md5.hexdigest()

def RunPreprocess(image_path, image_offset):
  """Preprocess an image and return back a collector and a fscache.

  Args:
    image_path: The path to the disk image to parse.
    image_offset: The sector offset into the volume where the partition starts.

  Returns:
    A tuple with three attributes, a collector, fileystem cache object and
    a pre processing object.
  """
  fscache = pfile.FilesystemCache()
  pre_obj = preprocess.PlasoPreprocess()

  # Start with a regular TSK collector.
  try:
    image_collector = collector_factory.GetImagePreprocessCollector(
        pre_obj, image_path, byte_offset=(image_offset * 512))
  except errors.UnableToOpenFilesystem as e:
    raise RuntimeError('Unable to proceed, not an image file? [%s]' % e)

  plugin_list = preprocessors.PreProcessList(pre_obj, image_collector)

  logging.info('Guessing OS')
  guessed_os = preprocess.GuessOS(image_collector)
  logging.info('OS: %s', guessed_os)
  logging.info('Running preprocess.')
  for weight in plugin_list.GetWeightList(guessed_os):
    for plugin in plugin_list.GetWeight(guessed_os, weight):
      try:
        plugin.Run()
      except errors.PreProcessFail as e:
        logging.warning(
            'Unable to run preprocessor: %s, reason: %s',
            plugin.__class__.__name__, e)

  logging.info('Preprocess done, saving files from image.')

  return image_collector, fscache, pre_obj

def RunExtraction(options, image_collector, pre_obj, fscache):
  """Run the extraction process.

  This method runs the file extraction process on the image and
  potentially on every VSS if that is wanted.

  Args:
    options: A configuration object, something like an argparse object.
    image_collector: An image collector object.
    pre_obj: Pre.
    fscache: A filesystem cache object.
  """
  # Save the regular files.
  FileSaver.calc_md5 = options.remove_duplicates
  ExtractFiles(image_collector, options.filter, options.path, fscache=fscache)

  # Go through VSS if desired.
  if options.vss:
    logging.info('Extracting files from VSS.')
    vss_numbers = vss.GetVssStoreCount(options.image, options.offset * 512)
    for store_nr in range(0, vss_numbers):
      logging.info('Extracting files from VSS %d/%d', store_nr + 1, vss_numbers)
      vss_col = collector_factory.GetImagePreprocessCollector(
          pre_obj, options.image, byte_offset=(options.offset * 512),
          vss_store_number=store_nr)
      FileSaver.prefix = 'vss_%d' % store_nr
      ExtractFiles(vss_col, options.filter, options.path, fscache)

def Main():
  """The main function, running the show."""
  arg_parser = argparse.ArgumentParser(
      description=('This is a simple collector designed to export files'
                   ' inside an image, both within a regular raw image as '
                   'well as inside a VSS. The tool uses a collection filter '
                   'that uses the same syntax as a targeted plaso filter.'),
      epilog='And that\'s how you export files, plaso style.')

  arg_parser.add_argument(
      '-v', '--vss', dest='vss', action='store_true', default=False,
      help='Indicate to the tool that we want to extract files from VSS too.')

  arg_parser.add_argument(
      '-o', '--offset', dest='offset', action='store', default=0, type=int,
      help='Offset in sectors into where the partition starts.')

  arg_parser.add_argument(
      '-w', '--write', dest='path', action='store', default='.', type=str,
      help='The directory in which extracted files should be stored in.')

  arg_parser.add_argument(
      '-x', '--extensions', dest='extension_string', action='store',
      type=str, metavar='EXTENSION_STRING', help=(
          'If the purpose is to find all files given a certain extension'
          'this options should be used. This option accepts a comma separated'
          'string denoting all file extensions, eg: -x "csv,docx,pst".'))

  arg_parser.add_argument(
      '-f', '--filter', action='store', dest='filter', metavar='FILTER_FILE',
      type=str, help=(
          'Full path to the file that contains the collection filter, '
          'the file can use variables that are defined in preprocesing'
          ', just like any other log2timeline/plaso collection filter.'))

  arg_parser.add_argument(
      '--include_duplicates', dest='include_duplicates', action='store_true',
      default=False, help=(
          'By default if VSS is turned on all files saved will have their'
          'MD5 sum calculated and compared to other files already saved '
          'with the same inode value. If the MD5 sum is the same the file '
          'does not get saved again. This option turns off that behavior '
          'so that all files will get stored, even if they are duplicates.'))

  arg_parser.add_argument(
      'image', action='store', metavar='IMAGE', default=None, type=str, help=(
          'The full path to the image file that we are about to extract files'
          ' from, it should be a raw image or another image that plaso '
          'supports.'))

  options = arg_parser.parse_args()

  format_str = '[%(levelname)s] %(message)s'
  logging.basicConfig(level=logging.INFO, format=format_str)

  if not os.path.isfile(options.image):
    raise RuntimeError('Unable to proceed, image file does not exist.')

  if not (options.filter or options.extension_string):
    arg_parser.print_help()
    print ''
    logging.error('Neither extension string nor filter defined.')
    sys.exit(1)

  if options.filter:
    if not os.path.isfile(options.filter):
      arg_parser.print_help()
      print ''
      logging.error('Unable to proceed, filter file does not exist.')
      sys.exit(1)

  if options.extension_string:
    extensions = [x.strip() for x in options.extension_string.split(',')]

  if not os.path.isdir(options.path):
    os.makedirs(options.path)

  if options.vss:
    if options.include_duplicates:
      options.remove_duplicates = False
    else:
      options.remove_duplicates = True
  else:
    options.remove_duplicates = False

  collector_use, fscache, pre_obj = RunPreprocess(options.image, options.offset)

  # The option of running both options.
  if options.filter:
    RunExtraction(options, collector_use, pre_obj, fscache)
    logging.info('Files from filters extracted.')

  if options.extension_string:
    RunExtensionExtraction(options, extensions, fscache)
    logging.info('Files based on extension extracted.')

if __name__ == '__main__':
  Main()
<commit_msg>Code review: 42520043:  PyVFS preparations refactor step 3.<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2013 The Plaso Project Authors.
# Please see the AUTHORS file for details on individual authors.
#
# 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.
"""A simple script to export files out of an image."""

import argparse
import hashlib
import logging
import os
import sys

from plaso import preprocessors

from plaso.collector import factory as collector_factory
from plaso.lib import collector_filter
from plaso.lib import errors
from plaso.lib import event
from plaso.lib import pfile
from plaso.lib import preprocess
from plaso.lib import putils
from plaso.lib import queue
from plaso.lib import vss

class ImageExtractor(object):
  """Class that implements the image extractor."""

  def __init__(self):
    """Initializes the image extractor object."""
    super(ImageExtractor, self).__init__()
    self._image_path = None
    self._image_offset = 0
    self._fscache = pfile.FilesystemCache()
    self._image_collector = None
    self._pre_obj = None

  def Open(self, image_path, sector_offset=0):
    """Opens the image.

    Args:
      image_path: The path to the disk image to parse.
      sector_offset: Option offset in sectors into the image where the partition
                     containing the file system starts. The default is 0.
    """
    self._image_path = image_path
    self._image_offset = sector_offset * 512

  def _Preprocess(self):
    """Preprocesses the image.

    Returns:
      The image collector object.
    """
    if self._pre_obj is not None:
      return

    self._pre_obj = preprocess.PlasoPreprocess()
    try:
      image_collector = collector_factory.GetImagePreprocessCollector(
          self._pre_obj, self._image_path, byte_offset=self._image_offset)

    except errors.UnableToOpenFilesystem as e:
      raise RuntimeError('Unable to proceed, not an image file? [%s]' % e)

    plugin_list = preprocessors.PreProcessList(
        self._pre_obj, image_collector)

    logging.info('Guessing OS')
    guessed_os = preprocess.GuessOS(image_collector)
    logging.info('OS: %s', guessed_os)

    logging.info('Running preprocess.')
    for weight in plugin_list.GetWeightList(guessed_os):
      for plugin in plugin_list.GetWeight(guessed_os, weight):
        try:
          plugin.Run()
        except errors.PreProcessFail as e:
          logging.warning(
              'Unable to run preprocessor: %s, reason: %s',
              plugin.__class__.__name__, e)

    logging.info('Preprocess done, saving files from image.')

    return image_collector

  def _ExtractFiles(self, image_collector, filter_expression, destination_path):
    """Extracts files that match the filter.

    Args:
      image_collector: The image collector object.
      filter_expression: String containing a collection filter expression.
      destination_path: The path where the extracted files should be stored.

    Returns:
      None, unless calc_md5 is set then it returns a dict with MD5 sums.
    """
    filter_object = collector_filter.CollectionFilter(
        image_collector, filter_expression)

    for path_spec in filter_object.GetPathSpecs():
      with pfile.OpenPFile(path_spec, fscache=self._fscache) as fh:
        # There will be issues on systems that use a different separator than a
        # forward slash. However a forward slash is always used in the pathspec.
        if os.path.sep != '/':
          fh.name = fh.name.replace('/', os.path.sep)
        FileSaver.SaveFile(fh, destination_path)

  def ExtractWithFilter(
      self, filter_expression, destination_path, process_vss=False,
      remove_duplicates=False):
    """Extracts files using a filter expression.

    This method runs the file extraction process on the image and
    potentially on every VSS if that is wanted.

    Args:
      filter_expression: String containing a collection filter expression.
      destination_path: The path where the extracted files should be stored.
      process_vss: Optional boolean value to indicate if VSS should be detected
                   and processed. The default is false.
      remove_duplicates: Optional boolean value to indicate if duplicates should
                         be detected and removed. The default is false.
    """
    if self._pre_obj is None:
      image_collector = self._Preprocess()

    if not os.path.isdir(destination_path):
      os.makedirs(destination_path)

    # Save the regular files.
    FileSaver.calc_md5 = remove_duplicates
    self._ExtractFiles(image_collector, filter_expression, destination_path)

    if process_vss:
      logging.info('Extracting files from VSS.')
      vss_numbers = vss.GetVssStoreCount(self._image_path, self._image_offset)

      for store_nr in range(0, vss_numbers):
        logging.info(
            'Extracting files from VSS %d/%d', store_nr + 1, vss_numbers)

        vss_collector = collector_factory.GetImagePreprocessCollector(
            self._pre_obj, self._image_path, byte_offset=self._image_offset,
            vss_store_number=store_nr)

        FileSaver.prefix = 'vss_%d' % store_nr
        self._ExtractFiles(vss_collector, filter_expression, destination_path)

  def ExtractWithExtensions(
       self, extensions, destination_path, process_vss=False,
       remove_duplicates=False):
    """Extracts files using extensions.

    Args:
      extensions: List of extensions.
      destination_path: The path where the extracted files should be stored.
      process_vss: Optional boolean value to indicate if VSS should be detected
                   and processed. The default is false.
      remove_duplicates: Optional boolean value to indicate if duplicates should
                         be detected and removed. The default is false.
    """
    if not os.path.isdir(destination_path):
      os.makedirs(destination_path)

    input_queue = queue.SingleThreadedQueue()
    output_queue = queue.SingleThreadedQueue()

    image_collector = collector_factory.GetImageCollector(
        input_queue, output_queue, self._image_path,
        byte_offset=self._image_offset, parse_vss=process_vss,
        fscache=self._fscache)

    image_collector.Collect()

    FileSaver.calc_md5 = remove_duplicates

    for path_spec_string in input_queue.PopItems():
      pathspec = event.EventPathSpec()
      pathspec.FromProtoString(path_spec_string)

      _, _, extension = pathspec.file_path.rpartition('.')
      if extension.lower() in extensions:
        fh = pfile.OpenPFile(pathspec, fscache=self._fscache)
        if getattr(pathspec, 'vss_store_number', None):
          FileSaver.prefix = 'vss_%d' % pathspec.vss_store_number + 1
        else:
          FileSaver.prefix = ''
        if os.path.sep != '/':
          fh.name = fh.name.replace('/', os.path.sep)

        FileSaver.SaveFile(fh, destination_path)

class FileSaver(object):
  """A simple class that is used to save files."""
  md5_dict = {}
  calc_md5 = False
  prefix = ''

  @classmethod
  def SaveFile(cls, fh, destination_path):
    """Take a filehandle and an export path and save the file."""
    directory = ''
    filename = ''
    if os.path.sep in fh.name:
      directory_string, _, filename = fh.name.rpartition(os.path.sep)
      if directory_string:
        directory = os.path.join(
            destination_path, *directory_string.split(os.path.sep))
    else:
      filename = fh.name

    if cls.prefix:
      extracted_filename = '{}_{}'.format(cls.prefix, filename)
    else:
      extracted_filename = filename

    while extracted_filename.startswith(os.path.sep):
      extracted_filename = extracted_filename[1:]

    if directory:
      if not os.path.isdir(directory):
        os.makedirs(directory)
    else:
      directory = destination_path

    save_file = True
    if cls.calc_md5:
      stat = fh.Stat()
      inode = stat.attributes.get('ino', 0)
      md5sum = CalculateHash(fh)
      if inode in cls.md5_dict:
        if md5sum in cls.md5_dict[inode]:
          save_file = False
        else:
          cls.md5_dict[inode].append(md5sum)
      else:
        cls.md5_dict[inode] = [md5sum]

    if save_file:
      try:
        putils.Pfile2File(fh, os.path.join(directory, extracted_filename))
      except IOError as e:
        logging.error(u'Unable to save file: {} [skipping]. Error {}'.format(
            fh.name, e))

def CalculateHash(file_object):
  """Return a hash for a given file object."""
  md5 = hashlib.md5()
  file_object.seek(0)

  data = file_object.read(4098)
  while data:
    md5.update(data)
    data = file_object.read(4098)

  return md5.hexdigest()

def Main():
  """The main function, running the show."""
  arg_parser = argparse.ArgumentParser(
      description=(
          'This is a simple collector designed to export files inside an '
          'image, both within a regular RAW image as well as inside a VSS. '
          'The tool uses a collection filter that uses the same syntax as a '
          'targeted plaso filter.'),
      epilog='And that\'s how you export files, plaso style.')

  arg_parser.add_argument(
      '-v', '--vss', dest='vss', action='store_true', default=False,
      help='Also extract files from VSS.')

  arg_parser.add_argument(
      '-o', '--offset', dest='offset', action='store', default=0, type=int,
      help='Offset in sectors into where the partition starts.')

  arg_parser.add_argument(
      '-w', '--write', dest='path', action='store', default='.', type=str,
      help='The directory in which extracted files should be stored in.')

  arg_parser.add_argument(
      '-x', '--extensions', dest='extension_string', action='store',
      type=str, metavar='EXTENSION_STRING', help=(
          'If the purpose is to find all files given a certain extension '
          'this options should be used. This option accepts a comma separated '
          'string denoting all file extensions, eg: -x "csv,docx,pst".'))

  arg_parser.add_argument(
      '-f', '--filter', action='store', dest='filter', metavar='FILTER_FILE',
      type=str, help=(
          'Full path to the file that contains the collection filter, '
          'the file can use variables that are defined in preprocesing, '
          'just like any other log2timeline/plaso collection filter.'))

  arg_parser.add_argument(
      '--include_duplicates', dest='include_duplicates', action='store_true',
      default=False, help=(
          'By default if VSS is turned on all files saved will have their '
          'MD5 sum calculated and compared to other files already saved '
          'with the same inode value. If the MD5 sum is the same the file '
          'does not get saved again. This option turns off that behavior '
          'so that all files will get stored, even if they are duplicates.'))

  arg_parser.add_argument(
      'image', action='store', metavar='IMAGE', default=None, type=str, help=(
          'The full path to the image file that we are about to extract files '
          'from, it should be a raw image or another image that plaso '
          'supports.'))

  options = arg_parser.parse_args()

  format_str = '[%(levelname)s] %(message)s'
  logging.basicConfig(level=logging.INFO, format=format_str)

  if not os.path.isfile(options.image):
    raise RuntimeError('Unable to proceed, image file does not exist.')

  if not (options.filter or options.extension_string):
    arg_parser.print_help()
    print ''
    logging.error('Neither extension string nor filter defined.')
    sys.exit(1)

  if options.filter:
    if not os.path.isfile(options.filter):
      arg_parser.print_help()
      print ''
      logging.error('Unable to proceed, filter file does not exist.')
      sys.exit(1)

  if not options.vss:
    options.remove_duplicates = False
  elif options.include_duplicates:
    options.remove_duplicates = False
  else:
    options.remove_duplicates = True

  image_extractor = ImageExtractor()
  image_extractor.Open(options.image, sector_offset=options.offset)

  # The option of running both options.
  if options.filter:
    image_extractor.ExtractWithFilter(
        options.path, options.filter,
        process_vss=options.vss,
        remove_duplicates=options.remove_duplicates)

  if options.extension_string:
    extensions = [x.strip() for x in options.extension_string.split(',')]

    image_extractor.ExtractWithExtensions(
        extensions, options.path, process_vss=options.vss,
        remove_duplicates=options.remove_duplicates)
    logging.info('Files based on extension extracted.')

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


<commit_before># -*- coding: utf-8 -*-
# for datasets acquired with topspin > 2.1 (requires INF)
# scale F1 dimension to get unified scale for MQMAS/STMAS experiments
# the program rewrites BF1/SFO1/O1/SW/SF/OFFSET according to
# the apparent Larmor frequency of isotropic dimension
# Select the nucleus to use for indirect dimension
# Select the MQMAS/STMAS type of experiment that gives the isotropic dimension
# option F2toF1 will read SF from F2 dimension (usefull for MQMAS/STMAS 2D)
# if F2 does not contain the same nucleus as F1 (MQ-INEPT experiment for example)
# then F1 must contain the original SFO1/SF data

from __future__ import division

import re
import os

def scalef1sheared(F2toF1=False, exptype='3QMAS', dataset=None):
    """
    Calculate the universal scale for MQMAS/STMAS experiments
    input args :
    F2toF1: (bool) if True uses SF from F2 dimension for calculating SF in F1
    exptype: (string) None (default), '3QMAS', 5QMAS' or 'STMAS': calculate new scale for exptype
           if None ask user for exptype
    dataset: the dataset as returned by CURDATA()
    """
    # if this function is called from imported module then
    # one needs to import TOPSPIN functions
    # so that they are available in the current namespace
    from TopCmds import CURDATA, RE, MSG, SELECT

    from JTutils.CpyLib import brukerPARIO
    if dataset == None:
        dataset = CURDATA()
    dtst = brukerPARIO.dataset(dataset)

    # check the dimensionality of dataset
    #maxdim = GETACQUDIM()
    #dimlist = range(1, maxdim+1)

    # initialize up to 4 channels
    m = [0, 0, 0, 0]
    NUCx = [0, 0, 0, 0]
    SFOx = [0, 0, 0, 0]
    Ox = [0, 0, 0, 0]
    BFx = [0, 0, 0, 0]
    boutons=[]
    mnemonique=[]
    channelmap=[]

    # find channels used in pp
    ppname = dtst.returnacqpath()+'/pulseprogram'
    ppfile = open(ppname,'r')
    pp = ppfile.read()
    ppfile.close()

    # check if channels f1,..,f4 are present in pp
    # channelmap lists the active channels
    for chan in range(0,4,1):
        NUCx[chan] = "not active"
        SFOx[chan] = "not active"
        Ox[chan] = "not active"
        BFx[chan] = "not active"
        matchpattern = "^[^;]*:f" + str(chan+1)
        m[chan] = re.search(matchpattern, pp, flags=re.M)
        if m[chan] :
            NUCx[chan] = dtst.readacqpar("NUC" + str(chan+1))
            if NUCx[chan] == "off":
                NUCx[chan] = "not active"
                continue
            SFOx[chan] = dtst.readacqpar("SFO" + str(chan+1))
            BFx[chan] = dtst.readacqpar("BF" + str(chan+1))
            Ox[chan] = dtst.readacqpar("O" + str(chan+1))
            mnemonique.append(str(chan+1))
            boutons.append("channel " + str(chan+1) + ": " + NUCx[chan])
            channelmap.append(chan)

    if F2toF1 :
        # F1 parameters are calculated from F2 (chan=0)
        chosenChan = 0
    else:
        # user selects the channel to set in F1 among the active channels
        channel = SELECT(title="available channels",
                         message="what channel do you want to use for F1 (click button) ?",
                         buttons=boutons, mnemonics=mnemonique)
        chosenChan = channelmap[channel]

    sfo = float(SFOx[chosenChan])
    bf = float(BFx[chosenChan])
    o = float(Ox[chosenChan])

    # get the spectral window / dwell time in F1
    swh = float(dtst.readacqpar("SW_h",dimension=2,status=True),)
    inf1 = 1e6/swh

    #MSG(str(swh)+"\n"+str(inf1))
    if F2toF1:
        sf = float(dtst.readprocpar("SF", dimension=1, status=False))
    else:
        sf = float(dtst.readprocpar("SF", dimension=2, status=False))

    # select spin number based on nucleus name from bruker table
    TOPSPIN_HOME = os.path.normpath(os.getenv("XWINNMRHOME"))
    nuclei = TOPSPIN_HOME + os.path.normpath("/exp/stan/nmr/lists/nuclei.all")
    spin = ''
    with open(nuclei,'r') as F:
        for line in F:
            if (line == ''):
                break
            line.strip()
            if (line.find('#') == 0):
                continue
            if (line.find(NUCx[chosenChan]) > 0):
                tmp=line.split()
                spinstr = tmp[4]
                break
        if spinstr[1] == '/' :
            spin = float(spinstr[0])/2
        else :
            spin = float(spinstr[0])

    theTypes = ['3QMAS', '5QMAS', 'STMAS']
    if exptype == None:
        typeIndex = SELECT(title="Kind of experiment",
                     message="What is the experiment that need universal scaling in F1?",
                     buttons=["3QMAS","5QMAS","STMAS"])
        exptype = theTypes[typeIndex]

    folding_times = "0"
    result = INPUT_DIALOG("number of times the spectrum is folded (-1, 0, +1,...)",
          """ Adjust the scale of a folded spectrum. Input how many times the spectrum is folded
          a positive number shift the scale to higher ppm value
          """,
             ["fold times="],
             [folding_times])
    folding_times = int(result[0])

    # calculate scaling factor based on experiment type and spin
    # 3QMAS correlates 3Q coherence with CT coherence
    if exptype == '3QMAS':
        ni = 3./2.
        mi = -3./2.
        nf = -1./2.
        mf = 1./2.
    # 5QMAS correlates 3Q coherence with CT coherence
    if exptype == '5QMAS':
        ni = 5./2.
        mi = -5./2.
        nf = -1./2.
        mf = 1./2.
    # STMAS correlates ST2 coherence with CT coherence
    if exptype == 'STMAS':
        ni = 3./2.
        mi = 1./2.
        nf = -1./2.
        mf = 1./2.

    mq = ni-mi
    I2 = spin*(spin + 1.)
    C4ni = ni*(18.*I2 - 34.*ni*ni - 5.)
    C4mi = mi*(18.*I2 - 34.*mi*mi - 5.)
    C4nf = nf*(18.*I2 - 34.*nf*nf - 5.)
    C4mf = mf*(18.*I2 - 34.*mf*mf - 5.)
    ratio = -(C4mi-C4ni)/(C4mf-C4nf)
    factor = abs(ratio - mq)
    #MSG("ratio R=" + str(ratio))

    #apply scaling factor to frequencies
    nuciso = "%s" % (NUCx[chosenChan])
    sf *= factor
    sfo *= factor
    bf *= factor
    o *= factor
# change the apparent carrier frequency if spectrum is folded
    sfo += folding_times*swh/1e6
    o += folding_times*swh
    # recalculate SW: note that topspin uses SFO instead of SF for calculating SW in ppm.
    # This make the SW parameter not accurately meaningful
    #
    newSW = swh/sfo

    # master parameters are BF,O,SW,SF and SW signification depend on
    # SFO : SW_h=SW*SFO
    # I'm not using Bruker function as relations make it difficult to predict
    # 1st : BF
    # then :O or SFO
    # finally: SW_h (that must be stored before hand) to recalculate SW
    # according to new new SFO
    # SF: seems to be fully independant

    # for scale to be OK after reprocessing: store modified acq parameters
    dtst.writeacqpar("NUC1", nuciso, dimension=2, status=True)
    dtst.writeacqpar("BF1", str(bf), dimension=2, status=True)
    dtst.writeacqpar("SFO1", str(sfo), dimension=2, status=True)
    dtst.writeacqpar("O1", str(o), dimension=2, status=True)
    dtst.writeacqpar("SW", str(newSW), dimension=2, status=True)
    # for current spectrum to look OK : stores modified proc parameters
    # from acqu and proc parameters
    si =  int(dtst.readprocpar("FTSIZE", dimension=2, status=True))
    stsi = int(dtst.readprocpar("STSI", dimension=2, status=True))
    stsr =  int(dtst.readprocpar("STSR", dimension=2, status=True))
    hzppt = swh/si
    swp = stsi*hzppt
    offset = ((sfo*1e6+swh/2)-stsr*hzppt -sf*1e6)/sf
    dtst.writeprocpar("SW_p", str(swp), dimension=2, status=True)
    dtst.writeprocpar("SF", str(sf), dimension=2, status=False)
    dtst.writeprocpar("SF", str(sf), dimension=2, status=True)
    dtst.writeprocpar("OFFSET",str(offset),dimension=2,status=True)

    RE(dataset)

if __name__ == '__main__':
    description = """
    Adjust the F1 scale for isotropic dimension in MQMAS/STMAS of MQHETCOR type experiments.
    """
    class dummy():
        def __init__(self):
            self.F2toF1 = False
            self.Q3 = False
            self.Q5 = False
            self.ST = False
    try :
        import argparse
        parser = argparse.ArgumentParser(description=description)
        parser.add_argument('--F2toF1', action='store_true',
                    help='use parameters from F2 dimension to calculate scale in F1 (ex.: for MQMAS)')
        group = parser.add_mutually_exclusive_group()
        group.add_argument('--3qmas', action='store_true', dest='Q3',
                    help='Calculate the scale for a 3QMAS experiment')
        group.add_argument('--5qmas', action='store_true', dest='Q5',
                    help='Calculate the scale for a 5QMAS experiment')
        group.add_argument('--stmas', action='store_true', dest='ST',
                    help='Calculate the scale for a STMAS experiment')
        args  =  parser.parse_args(sys.argv[1:])
    except ImportError:
        if len(sys.argv) > 1:
            MSG("Argparse module not found!\n Arguments won't be processed")
        args = dummy()
    except SystemExit:
        MSG(""" Script is exiting : either you asked for help or there is an argument error.
        Check console for additional information
        """  + parser.format_help() )
        EXIT()
    if args.Q3 : exptype = '3QMAS'
    elif args.Q5 : exptype = '5QMAS'
    elif args.ST : exptype = 'STMAS'
    else : exptype = None
    dataset = CURDATA()
    scalef1sheared(F2toF1=args.F2toF1, exptype = exptype, dataset=dataset)
<commit_msg>Read status SF. The non status doesn't seem to be used...<commit_after># -*- coding: utf-8 -*-
# for datasets acquired with topspin > 2.1 (requires INF)
# scale F1 dimension to get unified scale for MQMAS/STMAS experiments
# the program rewrites BF1/SFO1/O1/SW/SF/OFFSET according to
# the apparent Larmor frequency of isotropic dimension
# Select the nucleus to use for indirect dimension
# Select the MQMAS/STMAS type of experiment that gives the isotropic dimension
# option F2toF1 will read SF from F2 dimension (usefull for MQMAS/STMAS 2D)
# if F2 does not contain the same nucleus as F1 (MQ-INEPT experiment for example)
# then F1 must contain the original SFO1/SF data

from __future__ import division

import re
import os

def scalef1sheared(F2toF1=False, exptype='3QMAS', dataset=None):
    """
    Calculate the universal scale for MQMAS/STMAS experiments
    input args :
    F2toF1: (bool) if True uses SF from F2 dimension for calculating SF in F1
    exptype: (string) None (default), '3QMAS', 5QMAS' or 'STMAS': calculate new scale for exptype
           if None ask user for exptype
    dataset: the dataset as returned by CURDATA()
    """
    # if this function is called from imported module then
    # one needs to import TOPSPIN functions
    # so that they are available in the current namespace
    from TopCmds import CURDATA, RE, MSG, SELECT

    from JTutils.CpyLib import brukerPARIO
    if dataset == None:
        dataset = CURDATA()
    dtst = brukerPARIO.dataset(dataset)

    # check the dimensionality of dataset
    #maxdim = GETACQUDIM()
    #dimlist = range(1, maxdim+1)

    # initialize up to 4 channels
    m = [0, 0, 0, 0]
    NUCx = [0, 0, 0, 0]
    SFOx = [0, 0, 0, 0]
    Ox = [0, 0, 0, 0]
    BFx = [0, 0, 0, 0]
    boutons=[]
    mnemonique=[]
    channelmap=[]

    # find channels used in pp
    ppname = dtst.returnacqpath()+'/pulseprogram'
    ppfile = open(ppname,'r')
    pp = ppfile.read()
    ppfile.close()

    # check if channels f1,..,f4 are present in pp
    # channelmap lists the active channels
    for chan in range(0,4,1):
        NUCx[chan] = "not active"
        SFOx[chan] = "not active"
        Ox[chan] = "not active"
        BFx[chan] = "not active"
        matchpattern = "^[^;]*:f" + str(chan+1)
        m[chan] = re.search(matchpattern, pp, flags=re.M)
        if m[chan] :
            NUCx[chan] = dtst.readacqpar("NUC" + str(chan+1))
            if NUCx[chan] == "off":
                NUCx[chan] = "not active"
                continue
            SFOx[chan] = dtst.readacqpar("SFO" + str(chan+1))
            BFx[chan] = dtst.readacqpar("BF" + str(chan+1))
            Ox[chan] = dtst.readacqpar("O" + str(chan+1))
            mnemonique.append(str(chan+1))
            boutons.append("channel " + str(chan+1) + ": " + NUCx[chan])
            channelmap.append(chan)

    if F2toF1 :
        # F1 parameters are calculated from F2 (chan=0)
        chosenChan = 0
    else:
        # user selects the channel to set in F1 among the active channels
        channel = SELECT(title="available channels",
                         message="what channel do you want to use for F1 (click button) ?",
                         buttons=boutons, mnemonics=mnemonique)
        chosenChan = channelmap[channel]

    sfo = float(SFOx[chosenChan])
    bf = float(BFx[chosenChan])
    o = float(Ox[chosenChan])

    # get the spectral window / dwell time in F1
    swh = float(dtst.readacqpar("SW_h",dimension=2,status=True),)
    inf1 = 1e6/swh

    #MSG(str(swh)+"\n"+str(inf1))
    if F2toF1:
        # status of not status that is the question....
        sf = float(dtst.readprocpar("SF", dimension=1, status=True))
    else:
        sf = float(dtst.readprocpar("SF", dimension=2, status=True))
    MSG("sf=" + str(sf))
    # select spin number based on nucleus name from bruker table
    TOPSPIN_HOME = os.path.normpath(os.getenv("XWINNMRHOME"))
    nuclei = TOPSPIN_HOME + os.path.normpath("/exp/stan/nmr/lists/nuclei.all")
    spin = ''
    with open(nuclei,'r') as F:
        for line in F:
            if (line == ''):
                break
            line.strip()
            if (line.find('#') == 0):
                continue
            if (line.find(NUCx[chosenChan]) > 0):
                tmp=line.split()
                spinstr = tmp[4]
                break
        if spinstr[1] == '/' :
            spin = float(spinstr[0])/2
        else :
            spin = float(spinstr[0])

    theTypes = ['3QMAS', '5QMAS', 'STMAS']
    if exptype == None:
        typeIndex = SELECT(title="Kind of experiment",
                     message="What is the experiment that need universal scaling in F1?",
                     buttons=["3QMAS","5QMAS","STMAS"])
        exptype = theTypes[typeIndex]

    folding_times = "0"
    result = INPUT_DIALOG("number of times the spectrum is folded (-1, 0, +1,...)",
          """ Adjust the scale of a folded spectrum. Input how many times the spectrum is folded
          a positive number shift the scale to higher ppm value
          """,
             ["fold times="],
             [folding_times])
    folding_times = int(result[0])

    # calculate scaling factor based on experiment type and spin
    # 3QMAS correlates 3Q coherence with CT coherence
    if exptype == '3QMAS':
        ni = 3./2.
        mi = -3./2.
        nf = -1./2.
        mf = 1./2.
    # 5QMAS correlates 3Q coherence with CT coherence
    if exptype == '5QMAS':
        ni = 5./2.
        mi = -5./2.
        nf = -1./2.
        mf = 1./2.
    # STMAS correlates ST2 coherence with CT coherence
    if exptype == 'STMAS':
        ni = 3./2.
        mi = 1./2.
        nf = -1./2.
        mf = 1./2.

    mq = ni-mi
    I2 = spin*(spin + 1.)
    C4ni = ni*(18.*I2 - 34.*ni*ni - 5.)
    C4mi = mi*(18.*I2 - 34.*mi*mi - 5.)
    C4nf = nf*(18.*I2 - 34.*nf*nf - 5.)
    C4mf = mf*(18.*I2 - 34.*mf*mf - 5.)
    ratio = -(C4mi-C4ni)/(C4mf-C4nf)
    factor = abs(ratio - mq)
    #MSG("ratio R=" + str(ratio))

    #apply scaling factor to frequencies
    nuciso = "%s" % (NUCx[chosenChan])
    sf *= factor
    sfo *= factor
    bf *= factor
    o *= factor
# change the apparent carrier frequency if spectrum is folded
    sfo += folding_times*swh/1e6
    o += folding_times*swh
    # recalculate SW: note that topspin uses SFO instead of SF for calculating SW in ppm.
    # This make the SW parameter not accurately meaningful
    #
    newSW = swh/sfo

    # master parameters are BF,O,SW,SF and SW signification depend on
    # SFO : SW_h=SW*SFO
    # I'm not using Bruker function as relations make it difficult to predict
    # 1st : BF
    # then :O or SFO
    # finally: SW_h (that must be stored before hand) to recalculate SW
    # according to new new SFO
    # SF: seems to be fully independant

    # for scale to be OK after reprocessing: store modified acq parameters
    dtst.writeacqpar("NUC1", nuciso, dimension=2, status=True)
    dtst.writeacqpar("BF1", str(bf), dimension=2, status=True)
    dtst.writeacqpar("SFO1", str(sfo), dimension=2, status=True)
    dtst.writeacqpar("O1", str(o), dimension=2, status=True)
    dtst.writeacqpar("SW", str(newSW), dimension=2, status=True)
    # for current spectrum to look OK : stores modified proc parameters
    # from acqu and proc parameters
    si =  int(dtst.readprocpar("FTSIZE", dimension=2, status=True))
    stsi = int(dtst.readprocpar("STSI", dimension=2, status=True))
    stsr =  int(dtst.readprocpar("STSR", dimension=2, status=True))
    hzppt = swh/si
    swp = stsi*hzppt
    offset = ((sfo*1e6+swh/2)-stsr*hzppt -sf*1e6)/sf
    dtst.writeprocpar("SW_p", str(swp), dimension=2, status=True)
    dtst.writeprocpar("SF", str(sf), dimension=2, status=False)
    dtst.writeprocpar("SF", str(sf), dimension=2, status=True)
    dtst.writeprocpar("OFFSET",str(offset),dimension=2,status=True)

    RE(dataset)

if __name__ == '__main__':
    description = """
    Adjust the F1 scale for isotropic dimension in MQMAS/STMAS of MQHETCOR type experiments.
    """
    class dummy():
        def __init__(self):
            self.F2toF1 = False
            self.Q3 = False
            self.Q5 = False
            self.ST = False
    try :
        import argparse
        parser = argparse.ArgumentParser(description=description)
        parser.add_argument('--F2toF1', action='store_true',
                    help='use parameters from F2 dimension to calculate scale in F1 (ex.: for MQMAS)')
        group = parser.add_mutually_exclusive_group()
        group.add_argument('--3qmas', action='store_true', dest='Q3',
                    help='Calculate the scale for a 3QMAS experiment')
        group.add_argument('--5qmas', action='store_true', dest='Q5',
                    help='Calculate the scale for a 5QMAS experiment')
        group.add_argument('--stmas', action='store_true', dest='ST',
                    help='Calculate the scale for a STMAS experiment')
        args  =  parser.parse_args(sys.argv[1:])
    except ImportError:
        if len(sys.argv) > 1:
            MSG("Argparse module not found!\n Arguments won't be processed")
        args = dummy()
    except SystemExit:
        MSG(""" Script is exiting : either you asked for help or there is an argument error.
        Check console for additional information
        """  + parser.format_help() )
        EXIT()
    if args.Q3 : exptype = '3QMAS'
    elif args.Q5 : exptype = '5QMAS'
    elif args.ST : exptype = 'STMAS'
    else : exptype = None
    dataset = CURDATA()
    scalef1sheared(F2toF1=args.F2toF1, exptype = exptype, dataset=dataset)
<|endoftext|>