<commit_before>#!/usr/bin/env python
# Copyright 2011 The greplin-exception-catcher 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.

"""
Cron for sending exception logs to greplin-exception-catcher.

Usage: upload.py http://server.com secretKey exceptionDirectory
"""

import json
import os
import time
import os.path
import shutil
import stat
import sys
import urllib2
import signal

FILES_TO_KEEP = 2000
SETTINGS = {}

def reportTooManyExceptions(deleted, example):
  """Report a special message to the GEC server if we see too many exceptions."""
  newErr = {
    "project": "GEC",
    "environment": example["environment"],
    "serverName": example["serverName"],
    "message": "Too many exceptions - GEC deleted %d" % deleted,
    "timestamp": int(time.time())
  }
  sendException(newErr, "NoFilename")

def sendException(jsonData, filename):
  """Send an exception to the GEC server
     Returns True if sending succeeded
     If the send fails, returns False and moves the file to '_' + filename"""

  def handleError(e, message):
    """Handles an error by printing it and marking the file as available but with an additional strike."""
    print message % filename
    print e
    if filename != 'NoFilename':
      shutil.move(filename + '.processing', os.path.join(os.path.dirname(filename), '_' + os.path.basename(filename)))

  request = urllib2.Request('%s/report?key=%s' % (SETTINGS["server"], SETTINGS["secretKey"]),
                            json.dumps(jsonData),
                            {'Content-Type': 'application/json'})
  try:
    response = urllib2.urlopen(request)

  except urllib2.HTTPError, e:
    handleError(e, 'Error from server while uploading %s')
    print e.read()
    return False

  except urllib2.URLError, e:
    handleError(e, 'Error while uploading %s')
    return False

  status = response.getcode()

  if status != 200:
    raise Exception('Unexpected status code: %d' % status)

  return True

def processFiles(files):
  """Sends each exception file in files to gec"""

  # only keep the newest FILES_TO_KEEP entries
  outstanding = len(files)
  if outstanding > FILES_TO_KEEP:
    deleted = outstanding-FILES_TO_KEEP
    files = sorted(files, key=os.path.getctime)
    with open(files[0]) as f:
      reportTooManyExceptions(deleted, json.load(f))
    [os.remove(f) for f in files[0:deleted] if os.path.exists(f)]
    files = files[deleted:]

  for filename in files:
    if os.path.exists(filename):
      try:
        timestamp = os.stat(filename)[stat.ST_CTIME]
      except OSError:
        # Usually this happens when another process processes the same file.
        continue

      processingFilename = filename + '.processing'
      try:
        shutil.move(filename, processingFilename)
      except IOError:
        # Usually this happens when another process processes the same file.
        continue
      try:
        with open(processingFilename) as f:
          result = json.load(f)
      except ValueError, ex:
        with open(processingFilename) as f:
          print >> sys.stderr, "Could not read %s:" % filename
          print >> sys.stderr, '\n"""'
          print >> sys.stderr, f.read()
          print >> sys.stderr, '"""\n'
          print >> sys.stderr, str(ex)
        os.remove(processingFilename)
        continue
      result['timestamp'] = timestamp

      # Only delete on success - otherwise the file was moved
      if sendException(result, filename):
        os.remove(processingFilename)

def writePid(lockDir):
  """Writes the current pid to the lock directory."""
  with open(os.path.join(lockDir, 'pid'), 'w') as f:
    f.write(str(os.getpid()))

def main():
  """Runs the gec sender."""

  if len(sys.argv) not in (3, 4):
    print """USAGE: upload.py SERVER SECRET_KEY PATH [LOCKNAME]

LOCKNAME defaults to 'upload-lock'"""
    sys.exit(1)

  SETTINGS["server"] = sys.argv[1]
  SETTINGS["secretKey"] = sys.argv[2]
  path = sys.argv[3]

  lockName = sys.argv[4] if len(sys.argv) == 5 else 'upload-lock'
  lock = os.path.join(path, lockName)

  # mkdir will fail if the directory already exists, so we can use it as a file lock
  try:
    os.mkdir(lock)
    writePid(lock)
  except OSError:
    try:
      # Another upload.py instance is running! Kill it and take its place!
      with open(os.path.join(lock, 'pid')) as f:
        pid = int(f.read().strip())
      os.kill(pid, signal.SIGTERM)
      time.sleep(5)
      os.kill(pid, signal.SIGKILL)
      time.sleep(1)
      writePid(lock)
    except (IOError, ValueError, OSError):
      # No pid file, bad pid file, or pid doesn't exit. Fix this manually.
      print "Lock directory %r already exists." % lock
      print "Tried to find and kill an existing upload.py instance, but that didn't turn out right."
      print "Another upload.py appears to be running, maybe? Consider killing it and trying again."
      print "And remember to clean up the lock directory and the old pid file, if any."
      sys.exit(1)

  files = [os.path.join(path, f) for f in os.listdir(path)
           if f.endswith(".gec.json") and not '_____' in f]

  try:
    processFiles(files)
  finally:
    try:
      os.unlink(os.path.join(lock, 'pid'))
    except IOError:
      print >> sys.stderr, 'Tried to delete nonexistent pid file %r' % os.path.join(lock, 'pid')
      print >> sys.stderr, 'Lock directory %r existence status: %r' % os.path.exists(lock)
      if os.path.exists(lock):
        print >> sys.stderr, 'Lock directory contents: %r' % os.listdir(lock)
    try:
      os.rmdir(lock)
    except OSError:
      print >> sys.stderr, 'Could not delete lock directory %r (exists: %r)' % (lock, os.path.exists(lock))
      if os.path.exists(lock):
        print >> sys.stderr, 'Lock directory exists.'
        print >> sys.stderr, 'Lock directory contents: %r' % os.listdir(lock)

if __name__ == '__main__':
  main()
<commit_msg>If dying process takes lock directory with it, re-create the lockdir.<commit_after>#!/usr/bin/env python
# Copyright 2011 The greplin-exception-catcher 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.

"""
Cron for sending exception logs to greplin-exception-catcher.

Usage: upload.py http://server.com secretKey exceptionDirectory
"""

import json
import os
import time
import os.path
import shutil
import stat
import sys
import urllib2
import signal

FILES_TO_KEEP = 2000
SETTINGS = {}

def reportTooManyExceptions(deleted, example):
  """Report a special message to the GEC server if we see too many exceptions."""
  newErr = {
    "project": "GEC",
    "environment": example["environment"],
    "serverName": example["serverName"],
    "message": "Too many exceptions - GEC deleted %d" % deleted,
    "timestamp": int(time.time())
  }
  sendException(newErr, "NoFilename")

def sendException(jsonData, filename):
  """Send an exception to the GEC server
     Returns True if sending succeeded
     If the send fails, returns False and moves the file to '_' + filename"""

  def handleError(e, message):
    """Handles an error by printing it and marking the file as available but with an additional strike."""
    print message % filename
    print e
    if filename != 'NoFilename':
      shutil.move(filename + '.processing', os.path.join(os.path.dirname(filename), '_' + os.path.basename(filename)))

  request = urllib2.Request('%s/report?key=%s' % (SETTINGS["server"], SETTINGS["secretKey"]),
                            json.dumps(jsonData),
                            {'Content-Type': 'application/json'})
  try:
    response = urllib2.urlopen(request)

  except urllib2.HTTPError, e:
    handleError(e, 'Error from server while uploading %s')
    print e.read()
    return False

  except urllib2.URLError, e:
    handleError(e, 'Error while uploading %s')
    return False

  status = response.getcode()

  if status != 200:
    raise Exception('Unexpected status code: %d' % status)

  return True

def processFiles(files):
  """Sends each exception file in files to gec"""

  # only keep the newest FILES_TO_KEEP entries
  outstanding = len(files)
  if outstanding > FILES_TO_KEEP:
    deleted = outstanding-FILES_TO_KEEP
    files = sorted(files, key=os.path.getctime)
    with open(files[0]) as f:
      reportTooManyExceptions(deleted, json.load(f))
    [os.remove(f) for f in files[0:deleted] if os.path.exists(f)]
    files = files[deleted:]

  for filename in files:
    if os.path.exists(filename):
      try:
        timestamp = os.stat(filename)[stat.ST_CTIME]
      except OSError:
        # Usually this happens when another process processes the same file.
        continue

      processingFilename = filename + '.processing'
      try:
        shutil.move(filename, processingFilename)
      except IOError:
        # Usually this happens when another process processes the same file.
        continue
      try:
        with open(processingFilename) as f:
          result = json.load(f)
      except ValueError, ex:
        with open(processingFilename) as f:
          print >> sys.stderr, "Could not read %s:" % filename
          print >> sys.stderr, '\n"""'
          print >> sys.stderr, f.read()
          print >> sys.stderr, '"""\n'
          print >> sys.stderr, str(ex)
        os.remove(processingFilename)
        continue
      result['timestamp'] = timestamp

      # Only delete on success - otherwise the file was moved
      if sendException(result, filename):
        os.remove(processingFilename)

def writePid(lockDir):
  """Writes the current pid to the lock directory."""
  with open(os.path.join(lockDir, 'pid'), 'w') as f:
    f.write(str(os.getpid()))

def main():
  """Runs the gec sender."""

  if len(sys.argv) not in (3, 4):
    print """USAGE: upload.py SERVER SECRET_KEY PATH [LOCKNAME]

LOCKNAME defaults to 'upload-lock'"""
    sys.exit(1)

  SETTINGS["server"] = sys.argv[1]
  SETTINGS["secretKey"] = sys.argv[2]
  path = sys.argv[3]

  lockName = sys.argv[4] if len(sys.argv) == 5 else 'upload-lock'
  lock = os.path.join(path, lockName)

  # mkdir will fail if the directory already exists, so we can use it as a file lock
  try:
    os.mkdir(lock)
    writePid(lock)
  except OSError:
    try:
      # Another upload.py instance is running! Kill it and take its place!
      with open(os.path.join(lock, 'pid')) as f:
        pid = int(f.read().strip())
      os.kill(pid, signal.SIGTERM)
      time.sleep(5)
      os.kill(pid, signal.SIGKILL)
      time.sleep(1)
      try:
        os.mkdir(lock)
      except OSError:
        pass
      writePid(lock)
    except (IOError, ValueError, OSError):
      # No pid file, bad pid file, or pid doesn't exit. Fix this manually.
      print "Lock directory %r already exists." % lock
      print "Tried to find and kill an existing upload.py instance, but that didn't turn out right."
      print "Another upload.py appears to be running, maybe? Consider killing it and trying again."
      print "And remember to clean up the lock directory and the old pid file, if any."
      sys.exit(1)

  files = [os.path.join(path, f) for f in os.listdir(path)
           if f.endswith(".gec.json") and not '_____' in f]

  try:
    processFiles(files)
  finally:
    try:
      os.unlink(os.path.join(lock, 'pid'))
    except IOError:
      print >> sys.stderr, 'Tried to delete nonexistent pid file %r' % os.path.join(lock, 'pid')
      print >> sys.stderr, 'Lock directory %r existence status: %r' % os.path.exists(lock)
      if os.path.exists(lock):
        print >> sys.stderr, 'Lock directory contents: %r' % os.listdir(lock)
    try:
      os.rmdir(lock)
    except OSError:
      print >> sys.stderr, 'Could not delete lock directory %r (exists: %r)' % (lock, os.path.exists(lock))
      if os.path.exists(lock):
        print >> sys.stderr, 'Lock directory exists.'
        print >> sys.stderr, 'Lock directory contents: %r' % os.listdir(lock)

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


<commit_before>import errno
import io
import json
import logging
import os
import sys
import tempfile
import time
from datetime import datetime
from hashlib import md5
from optparse import OptionParser

from memento_damage.damage_analysis import MementoDamageAnalysis

base_dir = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, base_dir)
from memento_damage.tools import Command, rmdir_recursive, prompt_yes_no

class MementoDamage(object):
    # Define all output files
    APP_LOG_FILE_NAME = 'app.log'
    HTML_FILE_NAME = 'source.html'
    NETWORK_LOG_FILE_NAME = 'network.log'
    IMAGE_LOG_FILE_NAME = 'image.log'
    CSS_LOG_FILE_NAME = 'css.log'
    VIDEO_LOG_FILE_NAME = 'video.log'
    SCREENSHOT_FILE_NAME = 'screenshot.png'
    JSON_RESULT_FILE_NAME = 'result.json'

    background_color = 'FFFFFF'

    _crawljs_script = os.path.join(base_dir, 'phantomjs', 'crawl.js')
    _debug = False
    _info = False
    _mode = 'simple'
    _follow_redirection = False
    _clean_cache = True

    _result = None

    def __init__(self, uri, output_dir, options={}):
        self.uri = str(uri)
        self.output_dir = output_dir

        # Initialize variable
        self.app_log_file = os.path.join(self.output_dir, self.APP_LOG_FILE_NAME)
        self.html_file = os.path.join(self.output_dir, self.HTML_FILE_NAME)
        self.network_log_file = os.path.join(self.output_dir, self.NETWORK_LOG_FILE_NAME)
        self.image_log_file = os.path.join(self.output_dir, self.IMAGE_LOG_FILE_NAME)
        self.css_log_file = os.path.join(self.output_dir, self.CSS_LOG_FILE_NAME)
        self.video_log_file = os.path.join(self.output_dir, self.VIDEO_LOG_FILE_NAME)
        self.screenshot_file = os.path.join(self.output_dir, self.SCREENSHOT_FILE_NAME)
        self.json_result_file = os.path.join(self.output_dir, self.JSON_RESULT_FILE_NAME)

        # options
        if 'mode' in options: self._mode = options['mode']
        if 'redirect' in options: self._follow_redirection = options['redirect']
        if 'debug' in options:
            if options['debug'] == 'complete': self._debug = True
            elif options['debug'] == 'simple': self._info = True

        # Setup logger --> to show debug verbosity
        log_formatter = logging.Formatter('%(asctime)s %(name)s %(levelname)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')

        # To stdout
        log_stdout_handler = logging.StreamHandler(sys.stdout)
        log_stdout_handler.setFormatter(log_formatter)

        # To file
        log_file_handler = logging.FileHandler(self.app_log_file, mode='w')
        log_file_handler.setFormatter(log_formatter)
        log_file_handler.setLevel(logging.DEBUG)

        # Configure logger
        self.logger = logging.getLogger(uri)
        self.logger.addHandler(log_file_handler)
        self.logger.addHandler(log_stdout_handler)

        self.setup_logger()

    def setup_logger(self):
        if self._info:
            self.logger.setLevel(logging.INFO)
        elif self._debug:
            self.logger.setLevel(logging.DEBUG)
        else:
            self.logger.setLevel(logging.ERROR)

    def log_stdout(self, out, write_fn):
        if out and hasattr(out, 'readline'):
            for line in iter(out.readline, b''):
                line = line.strip()
                write_fn(line)
        else:
            write_fn(out)

    def log_output(self, msg):
        if 'background_color' in msg:
            msg = json.loads(msg)
            if msg['background_color']:
                self.background_color = msg['background_color']
                return

        if 'crawl-result' in msg:
            msg = json.loads(msg)
            crawl_result = msg['crawl-result']
            self._crawl_result = crawl_result
            msg = crawl_result['message']

        if self.logger.level == logging.DEBUG: self.logger.debug(msg)
        elif self.logger.level == logging.INFO: self.logger.info(msg)

    def log_stderr(self, out, write_fn):
        if out and hasattr(out, 'readline'):
            for line in iter(out.readline, b''):
                line = line.strip()
                write_fn(line)
        else:
            write_fn(out)

    def log_error(self, msg):
        if 'crawl-result' in msg:
            msg = json.loads(msg)
            crawl_result = msg['crawl-result']
            self._crawl_result = crawl_result

            if crawl_result['error']:
                self.response_time = datetime.now()

                if self._mode == 'simple':
                    self.logger.error(crawl_result['message'])
                elif self._mode == 'json':
                    self.logger.error(json.dumps(crawl_result, indent=4))
                else:
                    self.logger.error('Choose mode "simple" or "json"')
        else:
            self.logger.error(msg)

    def run(self):
        self.request_time = datetime.now()

        # Crawl page with phantomjs crawl.js via arguments
        # Equivalent with console:
        phantomjs = os.getenv('PHANTOMJS', 'phantomjs')

        pjs_cmd = [phantomjs, '--ssl-protocol=any', '--output-encoding=utf8', self._crawljs_script, self.uri,
                   self.output_dir, str(self._follow_redirection), str(self.logger.level)]
        cmd = Command(pjs_cmd, pipe_stdout_callback=self.log_stdout, pipe_stderr_callback=self.log_stderr)
        err_code = cmd.run(10 * 60,
                           stdout_callback_args=(self.log_output, ),
                           stderr_callback_args=(self.log_error, ))

        if err_code != 0:
            self._do_clean_cache()
            return

        # get result of damage analysis
        self._result = self._do_analysis()
        self.response_time = datetime.now()

        self._result['message'] = 'Calculation is finished in {} seconds'.format(
            (self.response_time - self.request_time).seconds)
        self._result['timer'] = {
            'request_time': (self.request_time - datetime(1970, 1, 1)).total_seconds(),
            'response_time': (self.response_time - datetime(1970, 1, 1)).total_seconds()
        }
        self._result['calculation_time'] = (self.response_time - self.request_time).seconds

        # Save output
        io.open(self.json_result_file, 'wb').write(json.dumps(self._result))

        self._do_clean_cache()
        return self._result

    def get_result(self):
        return self._result

    def print_result(self):
        # Print total damage
        if self._result:
            if self._mode == 'simple':
                print('Total damage of {} is {}'.format(self.uri, str(self._result['total_damage'])))
            elif self._mode == 'json':
                print(json.dumps(self._result, indent=4))
            else:
                self.logger.error('Choose mode "simple" or "json"')

    def _do_analysis(self):
        # Calculate damage
        analysis = MementoDamageAnalysis(self)
        analysis.run()

        return analysis.get_result()

    def _do_clean_cache(self):
        # Remove cache directory
        if self._clean_cache:
            time.sleep(3)
            rmdir_recursive(self.output_dir)

    def set_show_debug_message(self):
        self._debug = True
        self.setup_logger()

    def set_show_info_message(self):
        self._info = True
        self.setup_logger()

    def set_output_mode_json(self):
        self._mode = 'json'

    def set_follow_redirection(self):
        self._follow_redirection = True

    def set_dont_clean_cache_on_finish(self):
        self._clean_cache = False

def main():
    parser = OptionParser()
    parser.set_usage(parser.get_usage().replace('\n', '') + ' <URI>')
    parser.add_option("-O", "--output-dir",
                      dest="output_dir", default=None,
                      help="output directory (optional)")
    parser.add_option("--overwrite",
                      action="store_true", dest="overwrite", default=False,
                      help="overwrite existing output directory")
    parser.add_option("-m", "--mode",
                      dest="mode", default="simple",
                      help='output mode: "simple" or "json" [default: %default]')
    parser.add_option("-d", "--debug",
                      dest="debug", default="simple",
                      help='debug mode: "simple" or "complete" [default: %default]')
    parser.add_option("-L", "--redirect",
                      action="store_true", dest="redirect", default=False,
                      help="follow url redirection")

    (options, args) = parser.parse_args()
    options = vars(options)

    if len(args) < 1:
        parser.print_help()
        exit()

    uri = args[0]
    hashed_url = md5(uri).hexdigest()

    use_tempdir = False
    # If option -O is provided, use it
    if options['output_dir']:
        output_dir = options['output_dir']
        if not os.path.isabs(output_dir):
            output_dir = os.path.join(os.getcwd(), output_dir)
            output_dir = os.path.abspath(output_dir)

    # Otherwise make temp dir
    else:
        output_dir = tempfile.mkdtemp()
        use_tempdir = True

    output_dir = os.path.join(output_dir, hashed_url)

    # If overwrite is not provided in options --> Check whether output_dir is exists
    if not options['overwrite']:
        if os.path.exists(output_dir):
            overwrite = prompt_yes_no('Path "{}" is exists. Do yo want to overwrite?'.format(output_dir))
            # if not overwrite, dont continue
            if not overwrite:
                return

    # Make output_dir recursive
    try:
        os.makedirs(output_dir)
    except OSError, e:
        if e.errno != errno.EEXIST: raise

    # Instantiate and run
    damage = MementoDamage(uri, output_dir, options)
    if not use_tempdir:
        damage.set_dont_clean_cache_on_finish()
    damage.run()
    damage.print_result()

if __name__ == "__main__":
    main()<commit_msg>Change hashed_url to quoted_url<commit_after>import errno
import io
import json
import logging
import os
import sys
import tempfile
import time
import urllib
from datetime import datetime
from optparse import OptionParser

from memento_damage.damage_analysis import MementoDamageAnalysis

base_dir = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, base_dir)
from memento_damage.tools import Command, rmdir_recursive, prompt_yes_no

class MementoDamage(object):
    # Define all output files
    APP_LOG_FILE_NAME = 'app.log'
    HTML_FILE_NAME = 'source.html'
    NETWORK_LOG_FILE_NAME = 'network.log'
    IMAGE_LOG_FILE_NAME = 'image.log'
    CSS_LOG_FILE_NAME = 'css.log'
    VIDEO_LOG_FILE_NAME = 'video.log'
    SCREENSHOT_FILE_NAME = 'screenshot.png'
    JSON_RESULT_FILE_NAME = 'result.json'

    background_color = 'FFFFFF'

    _crawljs_script = os.path.join(base_dir, 'phantomjs', 'crawl.js')
    _debug = False
    _info = False
    _mode = 'simple'
    _follow_redirection = False
    _clean_cache = True

    _result = None

    def __init__(self, uri, output_dir, options={}):
        self.uri = str(uri)
        self.output_dir = output_dir

        # Initialize variable
        self.app_log_file = os.path.join(self.output_dir, self.APP_LOG_FILE_NAME)
        self.html_file = os.path.join(self.output_dir, self.HTML_FILE_NAME)
        self.network_log_file = os.path.join(self.output_dir, self.NETWORK_LOG_FILE_NAME)
        self.image_log_file = os.path.join(self.output_dir, self.IMAGE_LOG_FILE_NAME)
        self.css_log_file = os.path.join(self.output_dir, self.CSS_LOG_FILE_NAME)
        self.video_log_file = os.path.join(self.output_dir, self.VIDEO_LOG_FILE_NAME)
        self.screenshot_file = os.path.join(self.output_dir, self.SCREENSHOT_FILE_NAME)
        self.json_result_file = os.path.join(self.output_dir, self.JSON_RESULT_FILE_NAME)

        # options
        if 'mode' in options: self._mode = options['mode']
        if 'redirect' in options: self._follow_redirection = options['redirect']
        if 'debug' in options:
            if options['debug'] == 'complete': self._debug = True
            elif options['debug'] == 'simple': self._info = True

        # Setup logger --> to show debug verbosity
        log_formatter = logging.Formatter('%(asctime)s %(name)s %(levelname)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')

        # To stdout
        log_stdout_handler = logging.StreamHandler(sys.stdout)
        log_stdout_handler.setFormatter(log_formatter)

        # To file
        log_file_handler = logging.FileHandler(self.app_log_file, mode='w')
        log_file_handler.setFormatter(log_formatter)
        log_file_handler.setLevel(logging.DEBUG)

        # Configure logger
        self.logger = logging.getLogger(uri)
        self.logger.addHandler(log_file_handler)
        self.logger.addHandler(log_stdout_handler)

        self.setup_logger()

    def setup_logger(self):
        if self._info:
            self.logger.setLevel(logging.INFO)
        elif self._debug:
            self.logger.setLevel(logging.DEBUG)
        else:
            self.logger.setLevel(logging.ERROR)

    def log_stdout(self, out, write_fn):
        if out and hasattr(out, 'readline'):
            for line in iter(out.readline, b''):
                line = line.strip()
                write_fn(line)
        else:
            write_fn(out)

    def log_output(self, msg):
        if 'background_color' in msg:
            msg = json.loads(msg)
            if msg['background_color']:
                self.background_color = msg['background_color']
                return

        if 'crawl-result' in msg:
            msg = json.loads(msg)
            crawl_result = msg['crawl-result']
            self._crawl_result = crawl_result
            msg = crawl_result['message']

        if self.logger.level == logging.DEBUG: self.logger.debug(msg)
        elif self.logger.level == logging.INFO: self.logger.info(msg)

    def log_stderr(self, out, write_fn):
        if out and hasattr(out, 'readline'):
            for line in iter(out.readline, b''):
                line = line.strip()
                write_fn(line)
        else:
            write_fn(out)

    def log_error(self, msg):
        if 'crawl-result' in msg:
            msg = json.loads(msg)
            crawl_result = msg['crawl-result']
            self._crawl_result = crawl_result

            if crawl_result['error']:
                self.response_time = datetime.now()

                if self._mode == 'simple':
                    self.logger.error(crawl_result['message'])
                elif self._mode == 'json':
                    self.logger.error(json.dumps(crawl_result, indent=4))
                else:
                    self.logger.error('Choose mode "simple" or "json"')
        else:
            self.logger.error(msg)

    def run(self):
        self.request_time = datetime.now()

        # Crawl page with phantomjs crawl.js via arguments
        # Equivalent with console:
        phantomjs = os.getenv('PHANTOMJS', 'phantomjs')

        pjs_cmd = [phantomjs, '--ssl-protocol=any', '--output-encoding=utf8', self._crawljs_script, self.uri,
                   self.output_dir, str(self._follow_redirection), str(self.logger.level)]
        cmd = Command(pjs_cmd, pipe_stdout_callback=self.log_stdout, pipe_stderr_callback=self.log_stderr)
        err_code = cmd.run(10 * 60,
                           stdout_callback_args=(self.log_output, ),
                           stderr_callback_args=(self.log_error, ))

        if err_code != 0:
            self._do_clean_cache()
            return

        # get result of damage analysis
        self._result = self._do_analysis()
        self.response_time = datetime.now()

        self._result['message'] = 'Calculation is finished in {} seconds'.format(
            (self.response_time - self.request_time).seconds)
        self._result['timer'] = {
            'request_time': (self.request_time - datetime(1970, 1, 1)).total_seconds(),
            'response_time': (self.response_time - datetime(1970, 1, 1)).total_seconds()
        }
        self._result['calculation_time'] = (self.response_time - self.request_time).seconds

        # Save output
        io.open(self.json_result_file, 'wb').write(json.dumps(self._result))

        self._do_clean_cache()
        return self._result

    def get_result(self):
        return self._result

    def print_result(self):
        # Print total damage
        if self._result:
            if self._mode == 'simple':
                print('Total damage of {} is {}'.format(self.uri, str(self._result['total_damage'])))
            elif self._mode == 'json':
                print(json.dumps(self._result, indent=4))
            else:
                self.logger.error('Choose mode "simple" or "json"')

    def _do_analysis(self):
        # Calculate damage
        analysis = MementoDamageAnalysis(self)
        analysis.run()

        return analysis.get_result()

    def _do_clean_cache(self):
        # Remove cache directory
        if self._clean_cache:
            time.sleep(3)
            rmdir_recursive(self.output_dir)

    def set_show_debug_message(self):
        self._debug = True
        self.setup_logger()

    def set_show_info_message(self):
        self._info = True
        self.setup_logger()

    def set_output_mode_json(self):
        self._mode = 'json'

    def set_follow_redirection(self):
        self._follow_redirection = True

    def set_dont_clean_cache_on_finish(self):
        self._clean_cache = False

def main():
    parser = OptionParser()
    parser.set_usage(parser.get_usage().replace('\n', '') + ' <URI>')
    parser.add_option("-O", "--output-dir",
                      dest="output_dir", default=None,
                      help="output directory (optional)")
    parser.add_option("--overwrite",
                      action="store_true", dest="overwrite", default=False,
                      help="overwrite existing output directory")
    parser.add_option("-m", "--mode",
                      dest="mode", default="simple",
                      help='output mode: "simple" or "json" [default: %default]')
    parser.add_option("-d", "--debug",
                      dest="debug", default="simple",
                      help='debug mode: "simple" or "complete" [default: %default]')
    parser.add_option("-L", "--redirect",
                      action="store_true", dest="redirect", default=False,
                      help="follow url redirection")

    (options, args) = parser.parse_args()
    options = vars(options)

    if len(args) < 1:
        parser.print_help()
        exit()

    uri = args[0]
    quoted_url = urllib.quote(uri).replace('/', '_').replace('.', '-')

    use_tempdir = False
    # If option -O is provided, use it
    if options['output_dir']:
        output_dir = options['output_dir']
        if not os.path.isabs(output_dir):
            output_dir = os.path.join(os.getcwd(), output_dir)
            output_dir = os.path.abspath(output_dir)

    # Otherwise make temp dir
    else:
        output_dir = tempfile.mkdtemp()
        use_tempdir = True

    output_dir = os.path.join(output_dir, quoted_url)

    # If overwrite is not provided in options --> Check whether output_dir is exists
    if not options['overwrite']:
        if os.path.exists(output_dir):
            overwrite = prompt_yes_no('Path "{}" is exists. Do yo want to overwrite?'.format(output_dir))
            # if not overwrite, dont continue
            if not overwrite:
                return

    # Make output_dir recursive
    try:
        os.makedirs(output_dir)
    except OSError, e:
        if e.errno != errno.EEXIST: raise

    # Instantiate and run
    damage = MementoDamage(uri, output_dir, options)
    if not use_tempdir:
        damage.set_dont_clean_cache_on_finish()
    damage.run()
    damage.print_result()

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


<commit_before>import yaml
import os
import collections
from bcbio import utils
import subprocess
import functools
import stat
import difflib
from itertools import repeat
from functools import reduce

def in2out(in_file, word, transform=True, out_dir=None):
    """
    creates an out_file name from an in_file name by adding word to the
    in_file name and writing it to the output_directory if specified.
    if transform is True, the word replaces the extension of in_file.
    if False it adds it to the stem
    transform: in2out("test.bam", "sam", transform=True) -> "test.sam"
    non transform: in2out("test.bam", "sam", transform=False) -> "test_bam.sam"

    """
    (base, ext) = os.path.splitext(in_file)
    if out_dir:
        base = os.path.join(out_dir, os.path.basename(base))

    if transform:
        return "%s.%s" % (base, word)

    else:
        return "%s_%s.%s" % (base, ext, word)

def get_stem(filename):
    return os.path.basename(remove_suffix(filename))

def combine_pairs(input_files):
    """ calls files pairs if they are completely the same except
    for one has _1 and the other has _2 returns a list of tuples
    of pairs or singles """
    PAIR_FILE_IDENTIFIERS = ["1", "2"]

    pairs = []
    used = []
    for in_file in input_files:
        if in_file in used:
            continue
        for comp_file in input_files:
            if comp_file in used:
                continue
            s = difflib.SequenceMatcher(a=in_file, b=comp_file)
            blocks = s.get_matching_blocks()
            # length 3 means on match in the middle of the string
            if len(s.get_matching_blocks()) is not 3:
                continue
            if comp_file[blocks[0][2]] in PAIR_FILE_IDENTIFIERS:
                used.append(in_file)
                used.append(comp_file)
                pairs.append([in_file, comp_file])
                break
        if in_file not in used:
            pairs.append([in_file])
            used.append(in_file)

    return pairs

def freeze_files(files, directory):
    """ freezes a list of files, copying them to directory and
    setting them to be read only """
    new_files = [os.path.join(directory, x) for x in
                 map(os.path.basename, files)]
    [os.chmod(x, stat.S_IREAD | stat.S_IRGRP) for x in new_files]
    return new_files

# XXX I'm not 100% sure if this actually works, I forget use with caution
def memoize_outfile_to_dir(res_dir, ext):
    """Creates outfile from input file and ext, running if outfile not present.

    This requires a standard function usage. The first arg, or kwarg 'in_file', needs
    to be the input file that is being processed. The output name is created with the
    provided ext relative to the input. The function is only run if the created
    out_file is not present.
    Memoize_outfile was lifted from bcbio.utils (Brad Chapman), added the
    writing to a directory portion
    """
    utils.safe_makedir(res_dir)
    def decor(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            if len(args) > 0:
                in_file = args[0]
            else:
                in_file = kwargs['in_file']
            out_fname = "%s%s" % (os.path.splitext(in_file)[0], ext)
            out_file = os.path.join(res_dir, out_fname)
            if not os.path.exists(out_file) or os.path.getsize(out_file) == 0:
                kwargs['out_file'] = out_file
                f(*args, **kwargs)
            return out_file
        return wrapper
    return decor

def prepare_ref_file(ref, config):
    """Get a reference file, either by URL or locally.
    Lifted from Brad Chapman
    """
    url = ref.get("url", None)
    if url:
        ref_file = _download_ref(url, config["dir"]["ref"])
    else:
        ref_file = ref.get("file", None)
    assert ref_file is not None and os.path.exists(ref_file), ref_file
    return ref_file

def _download_ref(url, ref_dir):
    #Lifted from Brad Chapman
    dl_file = os.path.basename(url)
    ref_file = None
    for supported_ext, extract_cmd in [(".gz", "gunzip"),
                                       (".tgz", ("tar", "zxvf"))]:
        if dl_file.endswith(supported_ext):
            ref_file = os.path.join(ref_dir, dl_file[:-len(supported_ext)])
            break
    assert ref_file is not None, url
    if not os.path.exists(ref_file):
        with utils.chdir(ref_dir):
            cl = ["wget", url]
            subprocess.check_call(cl)
            cl = list(flatten([extract_cmd, dl_file]))
            subprocess.check_call(cl)
    return ref_file

def build_results_dir(stage_config, config):
    outdir = os.path.join(config["dir"]["results"],
                          stage_config["name"])
    return outdir

def transform_infile(filename, stage, delim="."):
    return replace_suffix(filename, stage.get("name", "transformed"), delim)

def filter_infile(filename, stage, delim="."):
    return append_stem(filename, stage.get("name", "filtered"), delim)

def remove_suffix(filename):
    filename, extension = os.path.splitext(filename)
    return filename

def append_stem(filename, word, delim="."):
    """ returns a filename with word appended to the stem
    example: /path/to/test.run.sam -> /path/to/test.run.sorted.sam """
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)
    fsplit.insert(len(fsplit) - 1, word)
    return os.path.join(dirname, delim.join(fsplit))

def replace_suffix(filename, suffix, delim="."):
    """ returns a filename with the suffix replaced
    with a new suffix
    exampoe: /path/to/test.run.sam -> /path/to/test.run.fa"""
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)[:-1]
    fsplit.append(suffix)
    return os.path.join(dirname, delim.join(fsplit))

def flatten_options(config):
    """ returns a list of pairs of options from a config file; returns
    an empty list if no options are found """
    return(flatten(config.get("options", {}.items())))

def flatten(l):
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
            for sub in flatten(el):
                yield sub
        else:
            yield el

def rfind_key(d, item):
    for k,v in d.items():
        if k == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_key(v, item)

def rfind_value(d, item):
    for k, v in d.items():
        if v == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_value(v, item)

def validate_config(config, data=True):
    # add the full path to any executables to avoid confusion later
    for k, v in config["program"].items():
        config["program"][k] = _add_full_path(v)

    if not all([os.path.exists(x) for x in config["program"].values()]):
        print("Cannot find all of the files or directories listed "
              "in 'program'.")
        exit(-1)

    if not all(x in get_stages(config) for x in config["run"]):
        print("Cannot find all of the stages listed to run.")
        exit(-1)

    if data and not all([os.path.exists(x) for x in config["input"]]):
        print("All of the input files do not exist.")
        exit(-1)

    return(config)

def get_stages(config):
    return config["stage"].keys()

def _add_full_path(program):
    """ add the full path to any executables """
    new_path = which(program)
    if new_path:
        program_config = new_path
    return program_config

def find_one_if(predicate, coll):
    """ returns a single entry in a collection if it contains exactly
    one true instance of a predicate and None otherwise """
    results = filter(predicate, coll)
    if len(results) != 1:
        return None
    else:
        print results
        return results[0]

def load_yaml(filename):
    with open(filename) as f:
        body = yaml.load(f)
        return body

def which(program):
    """ returns the path to an executable"""

    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

"""
simple file tracking to keep track of the transformations a file has made
just playing with it for now
"""

class FileWithHistory(object):
    """ this tracks transformations of an input file so we can get the
    metadata associated with it at any point """

    def __init__(self, in_file, **kwargs):
        self._name_list = [in_file]
        [setattr(self, key, value) for key, value in kwargs.items()]

    def curr_file(self):
        return self._name_list[len(self._name_list)]

    def add_file(self, in_file):
        self._name_list.append(in_file)

    def was_file(self, in_file):
        return in_file in self._name_list

def update_tracker(tracker, curr_files):
    new_tracker = [track.add_file(new_file) for (track, new_file)
                   in zip(tracker, curr_files)]
    return new_tracker

def curr_files(tracker):
    return [x.curr_file() for x in tracker]

def dict_to_vectors(d):
    """ convert a dictionary to two vectors. returns a tuple of
    (key, value). For example: {"a": [1, 2, 3], "b": [4, 5]} returns
    (["a" "a" "a" "b" "b"], [1 2 3 4 5])

    """

    ks = []
    vs = []
    for k, v in d.items():
        ks += list(repeat(k, len(v)))
        vs += v
    return (ks, vs)

class dotdict(dict):
    """
    access dictioanry items via dot syntax
    via: http://parand.com/say/index.php/2008/10/24/python-dot-notation-dictionary-access/:
    """

    def __getattr__(self, attr):
        return self.get(attr, None)
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

def nested_lookup(d, t):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer
    """
    return reduce(lambda d, t: d.get(t, {}), t, d)

def get_in(d, t, default=None):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer

    example: get_in({1: {2: 3}}, (1, 2)) -> 3
    example: get_in({1: {2: 3}}, (2, 3)) -> {}
    """
    result = reduce(lambda d, t: d.get(t, {}), t, d)
    if not result:
        return default
    else:
        return result

def is_sequence(arg):
    """
    check if 'arg' is a sequence

    example: arg([]) -> True
    example: arg("lol") -> False

    """
    return (not hasattr(arg, "strip") and
            hasattr(arg, "__getitem__") or
            hasattr(arg, "__iter__"))

def is_pair(arg):
    """
    check if 'arg' is a two-item sequence

    """
    return is_sequence(arg) and len(arg) == 2
<commit_msg>Added locate to find input files to utils.<commit_after>import yaml
import os
import collections
from bcbio import utils
import subprocess
import functools
import stat
import difflib
from itertools import repeat
from functools import reduce
import fnmatch

def in2out(in_file, word, transform=True, out_dir=None):
    """
    creates an out_file name from an in_file name by adding word to the
    in_file name and writing it to the output_directory if specified.
    if transform is True, the word replaces the extension of in_file.
    if False it adds it to the stem
    transform: in2out("test.bam", "sam", transform=True) -> "test.sam"
    non transform: in2out("test.bam", "sam", transform=False) -> "test_bam.sam"

    """
    (base, ext) = os.path.splitext(in_file)
    if out_dir:
        base = os.path.join(out_dir, os.path.basename(base))

    if transform:
        return "%s.%s" % (base, word)

    else:
        return "%s_%s.%s" % (base, ext, word)

def get_stem(filename):
    return os.path.basename(remove_suffix(filename))

def combine_pairs(input_files):
    """ calls files pairs if they are completely the same except
    for one has _1 and the other has _2 returns a list of tuples
    of pairs or singles """
    PAIR_FILE_IDENTIFIERS = ["1", "2"]

    pairs = []
    used = []
    for in_file in input_files:
        if in_file in used:
            continue
        for comp_file in input_files:
            if comp_file in used:
                continue
            s = difflib.SequenceMatcher(a=in_file, b=comp_file)
            blocks = s.get_matching_blocks()
            # length 3 means on match in the middle of the string
            if len(s.get_matching_blocks()) is not 3:
                continue
            if comp_file[blocks[0][2]] in PAIR_FILE_IDENTIFIERS:
                used.append(in_file)
                used.append(comp_file)
                pairs.append([in_file, comp_file])
                break
        if in_file not in used:
            pairs.append([in_file])
            used.append(in_file)

    return pairs

def freeze_files(files, directory):
    """ freezes a list of files, copying them to directory and
    setting them to be read only """
    new_files = [os.path.join(directory, x) for x in
                 map(os.path.basename, files)]
    [os.chmod(x, stat.S_IREAD | stat.S_IRGRP) for x in new_files]
    return new_files

# XXX I'm not 100% sure if this actually works, I forget use with caution
def memoize_outfile_to_dir(res_dir, ext):
    """Creates outfile from input file and ext, running if outfile not present.

    This requires a standard function usage. The first arg, or kwarg 'in_file', needs
    to be the input file that is being processed. The output name is created with the
    provided ext relative to the input. The function is only run if the created
    out_file is not present.
    Memoize_outfile was lifted from bcbio.utils (Brad Chapman), added the
    writing to a directory portion
    """
    utils.safe_makedir(res_dir)
    def decor(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            if len(args) > 0:
                in_file = args[0]
            else:
                in_file = kwargs['in_file']
            out_fname = "%s%s" % (os.path.splitext(in_file)[0], ext)
            out_file = os.path.join(res_dir, out_fname)
            if not os.path.exists(out_file) or os.path.getsize(out_file) == 0:
                kwargs['out_file'] = out_file
                f(*args, **kwargs)
            return out_file
        return wrapper
    return decor

def prepare_ref_file(ref, config):
    """Get a reference file, either by URL or locally.
    Lifted from Brad Chapman
    """
    url = ref.get("url", None)
    if url:
        ref_file = _download_ref(url, config["dir"]["ref"])
    else:
        ref_file = ref.get("file", None)
    assert ref_file is not None and os.path.exists(ref_file), ref_file
    return ref_file

def _download_ref(url, ref_dir):
    #Lifted from Brad Chapman
    dl_file = os.path.basename(url)
    ref_file = None
    for supported_ext, extract_cmd in [(".gz", "gunzip"),
                                       (".tgz", ("tar", "zxvf"))]:
        if dl_file.endswith(supported_ext):
            ref_file = os.path.join(ref_dir, dl_file[:-len(supported_ext)])
            break
    assert ref_file is not None, url
    if not os.path.exists(ref_file):
        with utils.chdir(ref_dir):
            cl = ["wget", url]
            subprocess.check_call(cl)
            cl = list(flatten([extract_cmd, dl_file]))
            subprocess.check_call(cl)
    return ref_file

def build_results_dir(stage_config, config):
    outdir = os.path.join(config["dir"]["results"],
                          stage_config["name"])
    return outdir

def transform_infile(filename, stage, delim="."):
    return replace_suffix(filename, stage.get("name", "transformed"), delim)

def filter_infile(filename, stage, delim="."):
    return append_stem(filename, stage.get("name", "filtered"), delim)

def remove_suffix(filename):
    filename, extension = os.path.splitext(filename)
    return filename

def append_stem(filename, word, delim="."):
    """ returns a filename with word appended to the stem
    example: /path/to/test.run.sam -> /path/to/test.run.sorted.sam """
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)
    fsplit.insert(len(fsplit) - 1, word)
    return os.path.join(dirname, delim.join(fsplit))

def replace_suffix(filename, suffix, delim="."):
    """ returns a filename with the suffix replaced
    with a new suffix
    exampoe: /path/to/test.run.sam -> /path/to/test.run.fa"""
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)[:-1]
    fsplit.append(suffix)
    return os.path.join(dirname, delim.join(fsplit))

def flatten_options(config):
    """ returns a list of pairs of options from a config file; returns
    an empty list if no options are found """
    return(flatten(config.get("options", {}.items())))

def flatten(l):
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
            for sub in flatten(el):
                yield sub
        else:
            yield el

def rfind_key(d, item):
    for k,v in d.items():
        if k == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_key(v, item)

def rfind_value(d, item):
    for k, v in d.items():
        if v == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_value(v, item)

def validate_config(config, data=True):
    # add the full path to any executables to avoid confusion later
    for k, v in config["program"].items():
        config["program"][k] = _add_full_path(v)

    if not all([os.path.exists(x) for x in config["program"].values()]):
        print("Cannot find all of the files or directories listed "
              "in 'program'.")
        exit(-1)

    if not all(x in get_stages(config) for x in config["run"]):
        print("Cannot find all of the stages listed to run.")
        exit(-1)

    if data and not all([os.path.exists(x) for x in config["input"]]):
        print("All of the input files do not exist.")
        exit(-1)

    return(config)

def get_stages(config):
    return config["stage"].keys()

def _add_full_path(program):
    """ add the full path to any executables """
    new_path = which(program)
    if new_path:
        program_config = new_path
    return program_config

def find_one_if(predicate, coll):
    """ returns a single entry in a collection if it contains exactly
    one true instance of a predicate and None otherwise """
    results = filter(predicate, coll)
    if len(results) != 1:
        return None
    else:
        print results
        return results[0]

def load_yaml(filename):
    with open(filename) as f:
        body = yaml.load(f)
        return body

def which(program):
    """ returns the path to an executable"""

    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

"""
simple file tracking to keep track of the transformations a file has made
just playing with it for now
"""

class FileWithHistory(object):
    """ this tracks transformations of an input file so we can get the
    metadata associated with it at any point """

    def __init__(self, in_file, **kwargs):
        self._name_list = [in_file]
        [setattr(self, key, value) for key, value in kwargs.items()]

    def curr_file(self):
        return self._name_list[len(self._name_list)]

    def add_file(self, in_file):
        self._name_list.append(in_file)

    def was_file(self, in_file):
        return in_file in self._name_list

def update_tracker(tracker, curr_files):
    new_tracker = [track.add_file(new_file) for (track, new_file)
                   in zip(tracker, curr_files)]
    return new_tracker

def curr_files(tracker):
    return [x.curr_file() for x in tracker]

def dict_to_vectors(d):
    """ convert a dictionary to two vectors. returns a tuple of
    (key, value). For example: {"a": [1, 2, 3], "b": [4, 5]} returns
    (["a" "a" "a" "b" "b"], [1 2 3 4 5])

    """

    ks = []
    vs = []
    for k, v in d.items():
        ks += list(repeat(k, len(v)))
        vs += v
    return (ks, vs)

class dotdict(dict):
    """
    access dictioanry items via dot syntax
    via: http://parand.com/say/index.php/2008/10/24/python-dot-notation-dictionary-access/:
    """

    def __getattr__(self, attr):
        return self.get(attr, None)
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

def nested_lookup(d, t):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer
    """
    return reduce(lambda d, t: d.get(t, {}), t, d)

def get_in(d, t, default=None):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer

    example: get_in({1: {2: 3}}, (1, 2)) -> 3
    example: get_in({1: {2: 3}}, (2, 3)) -> {}
    """
    result = reduce(lambda d, t: d.get(t, {}), t, d)
    if not result:
        return default
    else:
        return result

def is_sequence(arg):
    """
    check if 'arg' is a sequence

    example: arg([]) -> True
    example: arg("lol") -> False

    """
    return (not hasattr(arg, "strip") and
            hasattr(arg, "__getitem__") or
            hasattr(arg, "__iter__"))

def is_pair(arg):
    """
    check if 'arg' is a two-item sequence

    """
    return is_sequence(arg) and len(arg) == 2

def locate(pattern, root=os.curdir):
    '''Locate all files matching supplied filename pattern in and below
    supplied root directory.'''
    for path, dirs, files in os.walk(os.path.abspath(root)):
        for filename in fnmatch.filter(files, pattern):
            yield os.path.join(path, filename)
<|endoftext|>


<commit_before>import asyncio
import re
import datetime

import discord
import dateutil.parser

import utils

ARBITRARILY_LARGE_NUMBER = 10000000000000

class MessageCache:
   _SECRET_TOKEN = utils.SecretToken()

   _CH_JSON_FILENAME = "channel.json"

   @classmethod
   async def get_instance(cls, client, cache_directory):
      inst = cls(cls._SECRET_TOKEN)
      inst._client = client
      inst._data_dir = cache_directory + "messagecache/"
      inst._messages_before_initialization = []

      inst._data = {}
      # A tree of dictionaries in this arrangement:
      # server_id -> channel_id -> list of messages stored in tuples
      # Each message entry is a tuple.

      print("Caching messages...")
      print("This will take a while if a lot of messages are being read.")
      await inst._refresh_buffers()

      print("DEBUGGING MESSAGE CACHE:")
      for (serv_id, serv_dict) in inst._data.items():
         for (ch_id, ch_data) in serv_dict.items():
            ch = inst._client.search_for_channel(ch_id, enablenamesearch=False, serverrestriction=None)
            print("#" + ch.name + " has len " + str(len(ch_data)))
      print("DEBUGGING MESSAGE CACHE DONE!")
      return inst

   def __init__(self, token):
      if not token is self._SECRET_TOKEN:
         raise RuntimeError("Not allowed to instantiate directly. Please use get_instance().")
      return

   # PRECONDITION: Currently not operating on the buffers.
   async def record_message(self, msg):
      if not isinstance(msg.channel, discord.Channel):
         return # Never caches private messages.

      server_dict = None
      try:
         server_dict = self._data[msg.server.id]
      except KeyError:
         server_dict = {}
         self._data[msg.server.id] = server_dict

      ch_list = None
      try:
         ch_list = server_dict[msg.channel.id]
      except KeyError:
         ch_list = []
         server_dict[msg.channel.id] = ch_list

      ch_list.append(self._message_dict(msg))

      print("DEBUGGING MESSAGE CACHE:")
      for (serv_id, serv_dict) in self._data.items():
         for (ch_id, ch_data) in serv_dict.items():
            ch = self._client.search_for_channel(ch_id, enablenamesearch=False, serverrestriction=None)
            print("#" + ch.name + " has len " + str(len(ch_data)))
      print("DEBUGGING MESSAGE CACHE DONE!")

      return

   async def _refresh_buffers(self):
      await self._client.set_temp_game_status("Filling cache buffers.")
      for server in self._client.servers:
         server_dict = None
         try:
            server_dict = self._data[server.id]
         except KeyError:
            server_dict = {}
            self._data[server.id] = server_dict

         for ch in server.channels:
            if ch.type is discord.ChannelType.voice:
               continue
            print("FOUND CHANNEL " + ch.name)

            # TODO: Rename these variable names.
            ch_dir = self._get_ch_dir(server.id, ch.id)
            ch_json_filename = ch_dir + self._CH_JSON_FILENAME

            ch_json_data = None
            try:
               ch_json_data = utils.json_read(ch_json_filename)
            except FileNotFoundError:
               ch_json_data = {}

            ch_stored_timestamp = None
            try:
               ch_stored_timestamp = dateutil.parser.parse(ch_json_data["last message timestamp"])
            except KeyError:
               ch_stored_timestamp = datetime.datetime(datetime.MINYEAR, 1, 1)
               ch_json_data["last message timestamp"] = ch_stored_timestamp.isoformat()
               utils.json_write(ch_json_filename, data=ch_json_data)

            # This will now fill a buffer all messages of a channel.
            # TODO: Consider filling a file, then reading off the file.
            msg_buffer = []
            try:
               async for msg in self._client.logs_from(ch, limit=ARBITRARILY_LARGE_NUMBER):
                  if msg.timestamp <= ch_stored_timestamp:
                     break
                  # Insert in front since we're reading messages starting from most recent.
                  msg_buffer.insert(0, self._message_dict(msg))
            except discord.errors.Forbidden:
               pass

            server_dict[ch.id] = msg_buffer
      await self._client.remove_temp_game_status()
      return

   @classmethod
   def _message_dict(cls, msg):
      i = {}
      i["ti"] = msg.timestamp # Datetime
      i["id"] = msg.id # String
      i["au"] = msg.author.id # String
      i["co"] = msg.content # String
      i["at"] = msg.attachments # List of dictionaries.
      i["em"] = msg.embeds # List of dictionaries
      tmp = ""
      if not msg.edited_timestamp is None:
         tmp += "e" # for "edited"
      if msg.tts:
         tmp += "t"
      if msg.mention_everyone:
         tmp += "m"
      tmp += str(len(msg.mentions)) + "/" + str(len(msg.channel_mentions))
      i["fl"] = tmp
      return i

   def _get_ch_dir(self, server_id, ch_id):
      return self._data_dir + server_id + "/" + ch_id + "/"

<commit_msg>Successfully stores messages in files, mindful of the last message timestamp.<commit_after>import asyncio
import re
import datetime
import os

import discord
import dateutil.parser

import utils

ARBITRARILY_LARGE_NUMBER = 10000000000000

class MessageCache:
   _SECRET_TOKEN = utils.SecretToken()

   _CH_JSON_FILENAME = "channel.json"

   @classmethod
   async def get_instance(cls, client, cache_directory):
      inst = cls(cls._SECRET_TOKEN)
      inst._client = client
      inst._data_dir = cache_directory + "messagecache/"
      inst._messages_before_initialization = []

      inst._data = {}
      # A tree of dictionaries in this arrangement:
      # server_id -> channel_id -> list of messages stored in tuples
      # Each message entry is a tuple.

      print("Caching messages...")
      print("This will take a while if a lot of messages are being read.")
      await inst._fill_buffers()

      print("DEBUGGING MESSAGE CACHE:")
      for (serv_id, serv_dict) in inst._data.items():
         for (ch_id, ch_data) in serv_dict.items():
            ch = inst._client.search_for_channel(ch_id, enablenamesearch=False, serverrestriction=None)
            print("#" + ch.name + " has len " + str(len(ch_data)))
      print("DEBUGGING MESSAGE CACHE DONE!")
      return inst

   def __init__(self, token):
      if not token is self._SECRET_TOKEN:
         raise RuntimeError("Not allowed to instantiate directly. Please use get_instance().")
      return

   # PRECONDITION: Currently not operating on the buffers.
   async def record_message(self, msg):
      if not isinstance(msg.channel, discord.Channel):
         return # Never caches private messages.

      ch_dict = None
      try:
         ch_dict = self._data[msg.server.id]
      except KeyError:
         ch_dict = {}
         self._data[msg.server.id] = ch_dict

      ch_list = None
      try:
         ch_list = ch_dict[msg.channel.id]
      except KeyError:
         ch_list = []
         ch_dict[msg.channel.id] = ch_list

      ch_list.append(self._message_dict(msg))

      # Move to file if buffer is large enough.
      if len(self._data[msg.server.id][msg.channel.id]) >= 200:
         self._move_to_disk(msg.server.id, msg.channel.id)

      print("DEBUGGING MESSAGE CACHE:")
      for (serv_id, serv_dict) in self._data.items():
         for (ch_id, ch_data) in serv_dict.items():
            ch = self._client.search_for_channel(ch_id, enablenamesearch=False, serverrestriction=None)
            print("#" + ch.name + " has len " + str(len(ch_data)))
      print("DEBUGGING MESSAGE CACHE DONE!")

      return

   async def _fill_buffers(self):
      await self._client.set_temp_game_status("Filling cache buffers.")
      for server in self._client.servers:
         ch_dict = None
         try:
            ch_dict = self._data[server.id]
         except KeyError:
            ch_dict = {}
            self._data[server.id] = ch_dict

         for ch in server.channels:
            if ch.type is discord.ChannelType.voice:
               continue
            print("FOUND CHANNEL " + ch.name)

            # TODO: Rename these variable names.
            # TODO: Turn this into a function? (see duplicated code...)
            ch_dir = self._get_ch_dir(server.id, ch.id)
            ch_json_filepath = ch_dir + self._CH_JSON_FILENAME

            ch_json_data = None
            try:
               ch_json_data = utils.json_read(ch_json_filepath)
            except FileNotFoundError:
               ch_json_data = {}

            # TODO: Turn this into a function? (see duplicated code...)
            ch_stored_timestamp = None
            try:
               ch_stored_timestamp = dateutil.parser.parse(ch_json_data["last message timestamp"])
            except KeyError:
               ch_stored_timestamp = datetime.datetime(datetime.MINYEAR, 1, 1)
               ch_json_data["last message timestamp"] = ch_stored_timestamp.isoformat()
               utils.json_write(ch_json_filepath, data=ch_json_data)

            # This will now fill a buffer all messages of a channel.
            # TODO: Consider filling a file, then reading off the file.
            msg_buffer = []
            try:
               async for msg in self._client.logs_from(ch, limit=ARBITRARILY_LARGE_NUMBER):
                  if msg.timestamp <= ch_stored_timestamp:
                     break
                  # Insert in front since we're reading messages starting from most recent.
                  msg_buffer.insert(0, self._message_dict(msg))
            except discord.errors.Forbidden:
               pass

            ch_dict[ch.id] = msg_buffer

            # Move every 5000 messages to disk.
            while len(ch_dict[ch.id]) >= 5000:
               print("MOVING 5000 MESSAGES ON DISK. MSG_BUFFER LEN = " + str(len(ch_dict[ch.id])))
               self._move_to_disk(server.id, ch.id, messages=5000)

            # Move every 1000 messages to disk.
            while len(ch_dict[ch.id]) >= 1000:
               print("MOVING 1000 MESSAGES ON DISK. MSG_BUFFER LEN = " + str(len(ch_dict[ch.id])))
               self._move_to_disk(server.id, ch.id, messages=1000)

            # Now move every 200 messages to disk.
            while len(ch_dict[ch.id]) >= 200:
               print("MOVING 200 MESSAGES ON DISK. MSG_BUFFER LEN = " + str(len(ch_dict[ch.id])))
               self._move_to_disk(server.id, ch.id, messages=200)

      await self._client.remove_temp_game_status()
      return

   # Moves specified channel's messages to disk.
   # PARAMETER: messages - Number of earlier messages to move to a single file.
   #                       If None, then all messages are moved.
   # PRECONDITION: messages > 0
   # PRECONDITION: server_id and ch_id are both valid keys.
   def _move_to_disk(self, server_id, ch_id, messages=None):
      print("STORING MESSAGES ON DISK")
      ch_dir = self._get_ch_dir(server_id, ch_id)
      ch_json_filepath = ch_dir + self._CH_JSON_FILENAME

      # TODO: Turn this into a function? (see duplicated code...)
      ch_json_data = None
      try:
         ch_json_data = utils.json_read(ch_json_filepath)
      except FileNotFoundError:
         ch_json_data = {}

      # TODO: Turn this into a function? (see duplicated code...)
      ch_stored_timestamp = None
      try:
         ch_stored_timestamp = dateutil.parser.parse(ch_json_data["last message timestamp"])
      except KeyError:
         ch_stored_timestamp = datetime.datetime(datetime.MINYEAR, 1, 1)
         ch_json_data["last message timestamp"] = ch_stored_timestamp.isoformat()
         utils.json_write(ch_json_filepath, data=ch_json_data)

      ch_dict = self._data[server_id]
      print("DATA LEN = " + str(len(ch_dict[ch_id])))

      # Split off the messages to be stored.
      to_store = ch_dict[ch_id][:messages]
      ch_dict[ch_id] = ch_dict[ch_id][messages:]

      print("TOSTORE LEN = " + str(len(to_store)))
      print("NEW DATA LEN = " + str(len(ch_dict[ch_id])))

      latest_message = to_store[-1:][0]
      latest_timestamp_isoformat = latest_message["t"].isoformat()
      ch_json_data["last message timestamp"] = latest_timestamp_isoformat

      for msg_dict in to_store:
         msg_dict["t"] = msg_dict["t"].isoformat() # Make serializable

      # Check the highest numbered json file.
      highest_json_file_number = 0
      for file_name in os.listdir(ch_dir):
         if file_name.endswith(".json"):
            file_number = None
            try:
               file_number = int(file_name[:-5])
            except ValueError:
               continue
            if file_number > highest_json_file_number:
               highest_json_file_number = file_number

      # Store data in the next available json file number
      file_name = str(highest_json_file_number + 1) + ".json"
      utils.json_write(ch_dir + file_name, data=to_store)

      # Save latest message timestamp.
      utils.json_write(ch_json_filepath, data=ch_json_data)
      return

   @classmethod
   def _message_dict(cls, msg):
      i = {}
      i["t"] = msg.timestamp # Datetime (must convert before storing!!!)
      i["i"] = msg.id # String
      i["a"] = msg.author.id # String
      i["c"] = msg.content # String
      i["h"] = msg.attachments # List of dictionaries.
      i["e"] = msg.embeds # List of dictionaries
      tmp = ""
      if not msg.edited_timestamp is None:
         tmp += "e" # for "edited"
      if msg.tts:
         tmp += "t"
      if msg.mention_everyone:
         tmp += "m"
      tmp += str(len(msg.mentions)) + "/" + str(len(msg.channel_mentions))
      i["f"] = tmp
      return i

   def _get_ch_dir(self, server_id, ch_id):
      return self._data_dir + server_id + "/" + ch_id + "/"

<|endoftext|>


<commit_before># bl.emailer - a mailer class

import sys, smtplib, traceback
from email.mime.text import MIMEText
from bl.dict import Dict

class Emailer(Dict):
    """create and send email using templates. Parameters (*=required):
        template*       : must have a render() method that takes keyword arguments
                            (if loader_class and Email.template_path are both given,
                                template can be a relative path)
        loader_class    : the class used to load templates (defaults to tornado.template.Loader).
        Email           : parameters from the Email config.
            template_path: the filesystem location in which to search for templates
            host        : The address of the mail host
            port        : The port the host is listening on
            from_address: The default From address
            to_address  : The default To address
            delivery    : 'smtp' sends it, 'test' returns the rendered message
            username    : username for smtp auth
            password    : password for smtp auth
            debug       : Whether debugging is on
    """

    def __init__(self, template_path=None, loader_class=None, log=print, **Email):
        """set up an emailer with a particular Email config"""
        Dict.__init__(self, log=log, **Email)
        if template_path is not None and loader_class is not None:
            self.loader = loader_class(template_path)

    def __repr__(self):
        return "Emailer(%r)" % (self.__file__)

    def render(self, template, **context):
        """render the emailer template with the given context."""
        if type(template)==str and self.loader is not None:
            template = self.loader.load(template)
        try:
            r = template.render(**context)
        except:
            r = template.generate(**context)
        if type(r)==type(b''): r = r.decode('UTF-8')
        return r

    def message(self, template, to_addr=None, subject=None, from_addr=None, cc=None, bcc=None, **context):
        """create a MIMEText message from the msg text with the given msg args"""
        msg = MIMEText(self.render(template, to_addr=to_addr or self.to_address, from_addr=from_addr, cc=cc, bcc=bcc, **context))
        msg['From'] = from_addr or self.from_address
        msg['Subject'] = subject
        for addr in (to_addr or self.to_address or '').split(','):
            msg.add_header('To', addr.strip())
        for addr in (cc or '').split(','):
            msg.add_header('Cc', addr.strip())
        for addr in (bcc or '').split(','):
            msg.add_header('Bcc', addr.strip())
        return msg

    def send_message(self, template, to_addr=None, subject=None, from_addr=None, cc=None, bcc=None, **context):
        return self.send(
            self.message(template,
                to_addr=to_addr, from_addr=from_addr, cc=cc, bcc=bcc, subject=subject, **context))

    def send(self, msg):
        """send the given msg and return the status of the delivery.
        Returns None if delivery succeeded.
        Returns a sys.exc_info() tuple if the SMPT client raised an exception.
        """
        if self.delivery == 'test':
            # return the message as text that would be sent
            return msg.as_string()

        elif self.delivery == 'smtp':
            # parse the message and send it
            fromaddr = msg['From']
            tolist = (msg.get_all('To') or []) \
                    + (msg.get_all('Cc') or []) \
                    + (msg.get_all('Bcc') or [])
            try:
                if self.port is not None:
                    smtpclient = smtplib.SMTP(self.host or '127.0.0.1', self.port)
                else:
                    smtpclient = smtplib.SMTP(self.host or '127.0.0.1')
                smtpclient.set_debuglevel(self.DEBUG and 1 or 0)    # non-zero gives us exceptions when emailing.
                if self.username is not None and self.password is not None:
                    smtpclient.login(self.username, self.password)
                for toaddr in tolist:
                    smtpclient.sendmail(fromaddr, toaddr, msg.as_string())
                smtpclient.quit()
            except:
                if self.DEBUG==True:
                    self.log(traceback.format_exc())
                else:
                    self.log("Emailer exception:", sys.exc_info()[1], file=sys.stderr)
                return sys.exc_info()[1]        # return the exception rather than raising it -- the show must go on.
<commit_msg>template_path is indeed an attribute of the Email config, as originally envisioned<commit_after># bl.emailer - a mailer class

import sys, smtplib, traceback
from email.mime.text import MIMEText
from bl.dict import Dict

class Emailer(Dict):
    """create and send email using templates. Parameters (*=required):
        template*       : must have a render() method that takes keyword arguments
                            (if loader_class and Email.template_path are both given,
                                template can be a relative path)
        loader_class    : the class used to load templates (defaults to tornado.template.Loader).
        Email           : parameters from the Email config.
            template_path: the filesystem location in which to search for templates
            host        : The address of the mail host
            port        : The port the host is listening on
            from_address: The default From address
            to_address  : The default To address
            delivery    : 'smtp' sends it, 'test' returns the rendered message
            username    : username for smtp auth
            password    : password for smtp auth
            debug       : Whether debugging is on
    """

    def __init__(self, loader_class=None, log=print, **Email):
        """set up an emailer with a particular Email config"""
        Dict.__init__(self, log=log, **Email)
        if Email.template_path is not None and loader_class is not None:
            self.loader = loader_class(Email.template_path)

    def __repr__(self):
        return "Emailer(%r)" % (self.__file__)

    def render(self, template, **context):
        """render the emailer template with the given context."""
        if type(template)==str and self.loader is not None:
            template = self.loader.load(template)
        try:
            r = template.render(**context)
        except:
            r = template.generate(**context)
        if type(r)==type(b''): r = r.decode('UTF-8')
        return r

    def message(self, template, to_addr=None, subject=None, from_addr=None, cc=None, bcc=None, **context):
        """create a MIMEText message from the msg text with the given msg args"""
        msg = MIMEText(self.render(template, to_addr=to_addr or self.to_address, from_addr=from_addr, cc=cc, bcc=bcc, **context))
        msg['From'] = from_addr or self.from_address
        msg['Subject'] = subject
        for addr in (to_addr or self.to_address or '').split(','):
            msg.add_header('To', addr.strip())
        for addr in (cc or '').split(','):
            msg.add_header('Cc', addr.strip())
        for addr in (bcc or '').split(','):
            msg.add_header('Bcc', addr.strip())
        return msg

    def send_message(self, template, to_addr=None, subject=None, from_addr=None, cc=None, bcc=None, **context):
        return self.send(
            self.message(template,
                to_addr=to_addr, from_addr=from_addr, cc=cc, bcc=bcc, subject=subject, **context))

    def send(self, msg):
        """send the given msg and return the status of the delivery.
        Returns None if delivery succeeded.
        Returns a sys.exc_info() tuple if the SMPT client raised an exception.
        """
        if self.delivery == 'test':
            # return the message as text that would be sent
            return msg.as_string()

        elif self.delivery == 'smtp':
            # parse the message and send it
            fromaddr = msg['From']
            tolist = (msg.get_all('To') or []) \
                    + (msg.get_all('Cc') or []) \
                    + (msg.get_all('Bcc') or [])
            try:
                if self.port is not None:
                    smtpclient = smtplib.SMTP(self.host or '127.0.0.1', self.port)
                else:
                    smtpclient = smtplib.SMTP(self.host or '127.0.0.1')
                smtpclient.set_debuglevel(self.DEBUG and 1 or 0)    # non-zero gives us exceptions when emailing.
                if self.username is not None and self.password is not None:
                    smtpclient.login(self.username, self.password)
                for toaddr in tolist:
                    smtpclient.sendmail(fromaddr, toaddr, msg.as_string())
                smtpclient.quit()
            except:
                if self.DEBUG==True:
                    self.log(traceback.format_exc())
                else:
                    self.log("Emailer exception:", sys.exc_info()[1], file=sys.stderr)
                return sys.exc_info()[1]        # return the exception rather than raising it -- the show must go on.
<|endoftext|>