<commit_before># hangups client wrapper for h2x

import sys
import hangups
import time
import asyncio
from threading import Thread

class ClientWrapper:
	def __init__(self, h2x, user, userJID):
		self.h2x = h2x
		self.userdb = h2x.userdb
		self.user = user
		self.userJID = userJID

		self.thread = None
		self.loop = None

		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client wrapper created")

	# Provides token if refresh fails
	# Would be nice to ask user for token
	def getToken(self):
		return self.user.token

	def connect(self):
		print("Connect!!!")
		self.disconnect()

		self.thread = Thread(target = self.clientBody)
		self.thread.start()

	def disconnect(self):
		print("Disconnect!!!")

		if self.thread != None:
			print("Stopping client thread")

			future = asyncio.async(self.client.disconnect(), loop = self.loop)
			future.add_done_callback(lambda future: print("Disconnect done"))

			print("Trashing references")
			self.thread = None
			self.loop = None

		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client disconnected")

		print("Disconnect done")

	def clientBody(self):
		# Initialize asyncio loop for this thread
		self.loop = asyncio.new_event_loop()
		loop = self.loop
		asyncio.set_event_loop(loop)

		# Get outh2 cookies
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Getting cookies")
		cookies = hangups.auth.get_auth(self.getToken, self.userdb.tokenRefreshPath(self.user))

		# Create client
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Initializing client")
		self.client = hangups.Client(cookies)

		# Add state change observers
		self.client.on_connect.add_observer(self.onConnect)
		self.client.on_disconnect.add_observer(self.onDisconnect)
		self.client.on_reconnect.add_observer(self.onReconnect)

		# Connect and run client
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client connecting...")
		# This will return when connection ends
		try:
			print("Running in loop")
			loop.run_until_complete(self.client.connect())
			print("Loop done")
		finally:
			loop.close()

		# Notify about client termination
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client disconnected")
		print("Client thread terminates")

	@asyncio.coroutine
	def onConnect(self, initialData):
		print("Connected!")
		self.h2x.sendComponentPresence(self.userJID, "available", "Online")

		self.userList = yield from hangups.build_user_list(self.client, initialData)
		self.convList = hangups.ConversationList(self.client, initialData.conversation_states, self.userList, initialData.sync_timestamp)
		self.convList.on_event.add_observer(self.onEvent)

		# Send presence for users on contact list
		for user in self.userList.get_all():
			if user.is_self == False:
				self.h2x.sendPresence(self.hang2JID(user), self.userJID, "available", "Present in user list")

		print("Connection handler end")

	@asyncio.coroutine
	def onDisconnect(self):
		print("Disconnected")

	@asyncio.coroutine
	def onReconnect(self):
		print("Reconnected")

	def hang2JID(self, hangUser):
		return hangUser.id_.chat_id + "@" + self.h2x.config.JID

	def onEvent(self, convEvent):
		conv = self.convList.get(convEvent.conversation_id)
		user = conv.get_user(convEvent.user_id)
		print("Message")
		print("User:")
		print(vars(user))
		print("Conv:")
		print(vars(conv))

		for event in conv._events:
			print("Event")
			print(type(event))
			print(event.text)
			self.h2x.sendMessage(self.userJID, self.hang2JID(user), event.text)

<commit_msg>Stop deliverying all messages on every message received<commit_after># hangups client wrapper for h2x

import sys
import hangups
import time
import asyncio
from threading import Thread

class ClientWrapper:
	def __init__(self, h2x, user, userJID):
		self.h2x = h2x
		self.userdb = h2x.userdb
		self.user = user
		self.userJID = userJID

		self.thread = None
		self.loop = None

		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client wrapper created")

	# Provides token if refresh fails
	# Would be nice to ask user for token
	def getToken(self):
		return self.user.token

	def connect(self):
		print("Connect!!!")
		self.disconnect()

		self.thread = Thread(target = self.clientBody)
		self.thread.start()

	def disconnect(self):
		print("Disconnect!!!")

		if self.thread != None:
			print("Stopping client thread")

			future = asyncio.async(self.client.disconnect(), loop = self.loop)
			future.add_done_callback(lambda future: print("Disconnect done"))

			print("Trashing references")
			self.thread = None
			self.loop = None

		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client disconnected")

		print("Disconnect done")

	def clientBody(self):
		# Initialize asyncio loop for this thread
		self.loop = asyncio.new_event_loop()
		loop = self.loop
		asyncio.set_event_loop(loop)

		# Get outh2 cookies
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Getting cookies")
		cookies = hangups.auth.get_auth(self.getToken, self.userdb.tokenRefreshPath(self.user))

		# Create client
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Initializing client")
		self.client = hangups.Client(cookies)

		# Add state change observers
		self.client.on_connect.add_observer(self.onConnect)
		self.client.on_disconnect.add_observer(self.onDisconnect)
		self.client.on_reconnect.add_observer(self.onReconnect)

		# Connect and run client
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client connecting...")
		# This will return when connection ends
		try:
			print("Running in loop")
			loop.run_until_complete(self.client.connect())
			print("Loop done")
		finally:
			loop.close()

		# Notify about client termination
		self.h2x.sendComponentPresence(self.userJID, "unavailable", "Client disconnected")
		print("Client thread terminates")

	@asyncio.coroutine
	def onConnect(self, initialData):
		print("Connected!")
		self.h2x.sendComponentPresence(self.userJID, "available", "Online")

		self.userList = yield from hangups.build_user_list(self.client, initialData)
		self.convList = hangups.ConversationList(self.client, initialData.conversation_states, self.userList, initialData.sync_timestamp)
		self.convList.on_event.add_observer(self.onEvent)

		# Send presence for users on contact list
		for user in self.userList.get_all():
			if user.is_self == False:
				self.h2x.sendPresence(self.hang2JID(user), self.userJID, "available", "Present in user list")

		print("Connection handler end")

	@asyncio.coroutine
	def onDisconnect(self):
		print("Disconnected")

	@asyncio.coroutine
	def onReconnect(self):
		print("Reconnected")

	def hang2JID(self, hangUser):
		return hangUser.id_.chat_id + "@" + self.h2x.config.JID

	def onEvent(self, convEvent):
		conv = self.convList.get(convEvent.conversation_id)
		user = conv.get_user(convEvent.user_id)
		self.h2x.sendMessage(self.userJID, self.hang2JID(user), convEvent.text)

<|endoftext|>


<commit_before>#!flask/bin/python
from flask import Flask, jsonify
from flask import abort
from flask import request
from flask import url_for
from json_flask import make_json_app
import db

app = make_json_app(__name__)

@app.route('/')
def index():
    return 'refla-todo'

@app.route('/api')
def api_root():
    return jsonify({
        'tasks': url_for('get_tasks'),
    })

@app.route('/api/tasks', methods=['GET'])
def get_tasks():
    tasks = db.get_all()
    return jsonify({'tasks': tasks})

@app.route('/api/tasks/<int:task_id>', methods=['GET'])
def get_task(task_id):
    task = db.get(task_id)
    if task is None:
        abort(404)
    return jsonify({'task': task})

@app.route('/api/tasks', methods=['POST'])
def create_task():
    if not request.json or not 'title' in request.json:
        abort(400)
    task = db.insert({
        'title': request.json['title'],
        'done': False
    })
    return jsonify({'task': task}), 201

@app.route('/api/tasks/<int:task_id>', methods=['PUT'])
def update_task(task_id):
    task = db.get(task_id)
    if task is None:
        abort(404)
    if not request.json:
        abort(400)
    if 'title' in request.json and type(request.json['title']) != unicode:
        abort(400)
    if 'done' in request.json and type(request.json['done']) is not bool:
        abort(400)
    print(task)
    task = db.update(task_id, {
        'title': request.json.get('title', task['title']),
        'done': request.json.get('done', task['done'])
    })
    print(task)
    return jsonify({'task': task})

@app.route('/api/tasks/<int:task_id>', methods=['DELETE'])
def delete_task(task_id):
    task = db.get(task_id)
    if task is None:
        abort(404)
    db.delete(task_id)
    return jsonify({'result': True})

if __name__ == '__main__':
    app.run(debug=True)<commit_msg>make api root url_for handling external<commit_after>#!flask/bin/python
from flask import Flask, jsonify
from flask import abort
from flask import request
from flask import url_for
from json_flask import make_json_app
import db

app = make_json_app(__name__)

@app.route('/')
def index():
    return 'refla-todo'

@app.route('/api')
def api_root():
    return jsonify({
        'tasks': url_for('get_tasks', _external=True),
    })

@app.route('/api/tasks', methods=['GET'])
def get_tasks():
    tasks = db.get_all()
    return jsonify({'tasks': tasks})

@app.route('/api/tasks/<int:task_id>', methods=['GET'])
def get_task(task_id):
    task = db.get(task_id)
    if task is None:
        abort(404)
    return jsonify({'task': task})

@app.route('/api/tasks', methods=['POST'])
def create_task():
    if not request.json or not 'title' in request.json:
        abort(400)
    task = db.insert({
        'title': request.json['title'],
        'done': False
    })
    return jsonify({'task': task}), 201

@app.route('/api/tasks/<int:task_id>', methods=['PUT'])
def update_task(task_id):
    task = db.get(task_id)
    if task is None:
        abort(404)
    if not request.json:
        abort(400)
    if 'title' in request.json and type(request.json['title']) != unicode:
        abort(400)
    if 'done' in request.json and type(request.json['done']) is not bool:
        abort(400)
    print(task)
    task = db.update(task_id, {
        'title': request.json.get('title', task['title']),
        'done': request.json.get('done', task['done'])
    })
    print(task)
    return jsonify({'task': task})

@app.route('/api/tasks/<int:task_id>', methods=['DELETE'])
def delete_task(task_id):
    task = db.get(task_id)
    if task is None:
        abort(404)
    db.delete(task_id)
    return jsonify({'result': True})

if __name__ == '__main__':
    app.run(debug=True)<|endoftext|>


<commit_before># Client
# Sends batch of tasks to workers via webservice
__author__ = 'ITS_Student'

import random
import time
import os
import threading
import json
import sys
from requests import post
import base64

class Client:
    def __init__(self, workers):
        # List of address
        self.workers = workers

class Main(object):
    def __init__(self, workers_list, path_to_input, path_to_output, pack_quantity):
        self.workers_list = workers_list
        self.path_to_input = path_to_input

        output_length = len(path_to_output)
        if path_to_output[output_length - 1] == '/':
            self.path_to_output = path_to_output
        else:
            self.path_to_output = path_to_output + '/'

        self.pack_quantity = int(pack_quantity)
        self.file_list = os.listdir(self.path_to_input)
        self.node_list = []
        self.threads = []
        self.elapsed_time = 0

    def worker(self, node, file_list):
        for filename in file_list:
            print node.workers + ": Processing " + filename
            file_format = filename[-4:]
            open_file = open(self.path_to_input + filename, 'rb')
            read_file = open_file.read()
            open_file.close()
            encoded = base64.b64encode(read_file)

            #Change
            data = {
                    "data": encoded,
                    "format": file_format
                   }
            retval = post('http://'+ node.workers + '/api/imageconverter/v0.1/', data=data).json()
            result = base64.b64decode(retval['result'])

            file_open = open(self.path_to_output + filename, 'wb')
            file_open.write(result)
            file_open.close()
            #Change

    def run(self):
        if self.path_to_input[-1] != '/':
            self.path_to_input += '/'

        # Connect to server
        for workers in self.workers_list:
                self.node_list.append(Client(workers))

        i = 0
        j = self.pack_quantity
        server_selector = 0

        start_time = time.time()
        while True:
            if len(self.file_list) == 0:
                break

            if i == len(self.file_list):
                break

            if j > len(self.file_list):
                j = len(self.file_list)

            if server_selector == len(self.node_list):
                server_selector = 0

            t = threading.Thread(target=self.worker, args=(self.node_list[server_selector], self.file_list[i:j]))
            self.threads.append(t)
            t.start()
            for files in range(i, j):
                del self.file_list[0]

            server_selector += 1

        for thread in self.threads:
            thread.join()

        self.elapsed_time = time.time() - start_time

# Main Program
if __name__ == "__main__":
    path_to_input = None
    path_to_output = None

    print "How many servers do you want to use?"
    qty = raw_input()
    address_list = []
    for node in range(int(qty)):
        if node == 0:
            print "Usage: [ip_address:port]"
            print "Example: Localhost:9000"

        address_list.append(raw_input())

    print "Enter input folder"
    print "Usage: /path/to/input"
    print "Example: /home/input"
    while True:
        path_to_input = raw_input()
        if not os.path.isdir(path_to_input):
            print "Invalid argument: Input path is incorrect. Make sure it is a directory, not a file."

        else:
            break

    print "Enter outout folder"
    print "Usage: /path/to/output"
    print "Example: /home/output"
    while True:
        path_to_output = raw_input()
        if not os.path.isdir(path_to_output):
            print "Invalid argument: Input path is incorrect. Make sure it is a directory, not a file."

        else:
            break

    print "How many files to distribute ?"
    pack_quantity = raw_input()

    main = Main(address_list, path_to_input, path_to_output, pack_quantity)
    main.run()

    print "Sending time " + str(main.elapsed_time)

<commit_msg>Fixing typo LOL<commit_after># Client
# Sends batch of tasks to workers via webservice
__author__ = 'ITS_Student'

import random
import time
import os
import threading
import json
import sys
from requests import post
import base64

class Client:
    def __init__(self, workers):
        # List of address
        self.workers = workers

class Main(object):
    def __init__(self, workers_list, path_to_input, path_to_output, pack_quantity):
        self.workers_list = workers_list
        self.path_to_input = path_to_input

        output_length = len(path_to_output)
        if path_to_output[output_length - 1] == '/':
            self.path_to_output = path_to_output
        else:
            self.path_to_output = path_to_output + '/'

        self.pack_quantity = int(pack_quantity)
        self.file_list = os.listdir(self.path_to_input)
        self.node_list = []
        self.threads = []
        self.elapsed_time = 0

    def worker(self, node, file_list):
        for filename in file_list:
            print node.workers + ": Processing " + filename
            file_format = filename[-4:]
            open_file = open(self.path_to_input + filename, 'rb')
            read_file = open_file.read()
            open_file.close()
            encoded = base64.b64encode(read_file)

            #Change
            data = {
                    "data": encoded,
                    "format": file_format
                   }
            retval = post('http://'+ node.workers + '/api/imageconverter/v0.1/', data=data).json()
            result = base64.b64decode(retval['result'])

            file_open = open(self.path_to_output + filename, 'wb')
            file_open.write(result)
            file_open.close()
            #Change

    def run(self):
        if self.path_to_input[-1] != '/':
            self.path_to_input += '/'

        # Connect to server
        for workers in self.workers_list:
                self.node_list.append(Client(workers))

        i = 0
        j = self.pack_quantity
        server_selector = 0

        start_time = time.time()
        while True:
            if len(self.file_list) == 0:
                break

            if i == len(self.file_list):
                break

            if j > len(self.file_list):
                j = len(self.file_list)

            if server_selector == len(self.node_list):
                server_selector = 0

            t = threading.Thread(target=self.worker, args=(self.node_list[server_selector], self.file_list[i:j]))
            self.threads.append(t)
            t.start()
            for files in range(i, j):
                del self.file_list[0]

            server_selector += 1

        for thread in self.threads:
            thread.join()

        self.elapsed_time = time.time() - start_time

# Main Program
if __name__ == "__main__":
    path_to_input = None
    path_to_output = None

    print "How many servers do you want to use?"
    qty = raw_input()
    address_list = []
    for node in range(int(qty)):
        if node == 0:
            print "Usage: [ip_address:port]"
            print "Example: Localhost:9000"

        address_list.append(raw_input())

    print "Enter input folder"
    print "Usage: /path/to/input"
    print "Example: /home/input"
    while True:
        path_to_input = raw_input()
        if not os.path.isdir(path_to_input):
            print "Invalid argument: Input path is incorrect. Make sure it is a directory, not a file."

        else:
            break

    print "Enter output folder"
    print "Usage: /path/to/output"
    print "Example: /home/output"
    while True:
        path_to_output = raw_input()
        if not os.path.isdir(path_to_output):
            print "Invalid argument: Input path is incorrect. Make sure it is a directory, not a file."

        else:
            break

    print "How many files to distribute ?"
    pack_quantity = raw_input()

    main = Main(address_list, path_to_input, path_to_output, pack_quantity)
    main.run()

    print "Sending time " + str(main.elapsed_time)

<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2015 LCI Technology Group, LLC
# All rights reserved

import requests
import pysodium
import os
import base64
import json
import cmd

# CONSTANTS
HOMEDIR = os.path.expanduser("~")
ZKMDIR = os.path.join(HOMEDIR, '.zkm')
CONFIG = os.path.join(ZKMDIR, 'config')
CONTACT = os.path.join(ZKMDIR, 'contacts')

def load_json_file(filename):
    """
    Load a json file into a dictionary and return it.
    """
    with open(filename, 'rb') as f:
        data = f.read()

    try:
        return json.loads(data)
    except:
        return {}

def save_json_data(filename, data):
    """
    Save the given data to a file in JSON format.
    """
    with open(filename, 'wb') as f:
        f.write('{0}\n'.format(json.dumps(data)))

def encrypt(our_secret, our_public, their_public, msg):
    """
    Encrypt a message using the provided information.
    """
    their_public = base64.b64decode(their_public)
    our_secret = base64.b64decode(our_secret)
    nonce = pysodium.randombytes(pysodium.crypto_box_NONCEBYTES)
    enc = pysodium.crypto_box(msg, nonce, their_public, our_secret)

    nonce = base64.b64encode(nonce)
    enc = base64.b64encode(enc)

    # Return our public_key, nonce, and the encrypted message
    return ':'.join(our_public, nonce, enc)

def decrypt(our_secret, msg):
    """
    Decrypt a message using the provided information.
    """
    our_secret = base64.b64decode(our_secret)
    their_public, nonce, enc_msg = msg.split(':')
    their_public = base64.b64decode(their_public)
    nonce = base64.b64decode(nonce)
    enc_msg = base64.b64decode(enc_msg)
    dec_msg = pysodium.crypto_box_open(enc_msg, nonce, their_public, our_secret)

    # Return the sender's public key and the decrypted message.
    return their_public, dec_msg

def print_msg(contacts, their_public, msg):
    """
    Print the sender and message.

    Lookup the public key of the sender to see if they are in our
    contacts. If they are print the username, if not print the public key.
    """
    for username, contact_public in contacts.iteritems():
        if their_public == contact_public:
            sender = username
        else:
            sender = their_public

    print('{0}: {1}'.format(sender, msg))

def send(server, method, endpoint, data=None):
    """
    Send a message to the server and process the response.
    """
    url = '{0}{1}'.format(server, endpoint)

    if method == 'POST':
        resp = requests.post(url, data=data)

    if method == 'GET':
        resp = requests.get(url, params=data)

    if resp.status_code == 200:
        j = resp.json()
        if j['error'] is not None:
            print('[-] {0}'.format(j['error']))
            return None
        else:
            return j['response']
    else:
        print('[-] Server error: {0}'.format(resp.status_code))

def initialize():
    """
    Create a ~/.zkm directory with a config file inside.

    The config file will hold our public key, secret key, and since value.
    """
    if os.path.exists(ZKMDIR) is False:
        print('[+] Creating ZKM configuration directory.')
        os.mkdir(ZKMDIR, 0o750)

        print('[+] Creating new keypair.')
        our_public, our_secret = pysodium.crypto_sign_keypair()

        print('[+] Creating configuration file.')
        config = {'public': base64.b64encode(our_public),
                  'secret': base64.b64encode(our_secret),
                  'since': 1}

        save_json_data(CONFIG, config)
        os.chmod(CONFIG, 0o600)

        print('[+] Creating contacts file.')
        save_json_data(CONTACT, {})
        os.chmod(CONTACT, 0o600)

    else:
        print('[-] ZKM configuration directory already exists.')

class ZKMClient(cmd.Cmd):
    """
    ZKM: A zero knowledge messaging system.
    """
    prompt = 'zkm> '

    # Functions used in the interactive command prompt.
    def preloop(self):
        """
        Initialize the ZKM client if necessary.
        """
        try:
            self.config = load_json_file(CONFIG)
        except:
            print('[-] ZKM not initialized yet.')
            initialize()

        try:
            self.contacts = load_json_file(CONTACT)
        except Exception:
            print('[-] Could not load contacts file.')
            self.contacts = []

    def postloop(self):
        save_json_data(CONFIG, self.config)
        save_json_data(CONTACT, self.contacts)

    def do_add_contact(self, line):
        """
        Add a new contact to the contact list.
        """
        name, their_public = line.split(' ')
        self.contacts[name] = their_public
        save_json_data(CONTACT, self.contacts)

    def do_del_contact(self, name):
        """
        Remove a contact from the contact list.
        """
        self.contacts.pop(name, None)
        save_json_data(CONTACT, self.contacts)

    def do_connect(self, line):
        """
        Define the server we want to connect to for messages.
        """
        self.config['server'] = line
        save_json_data(CONFIG, self.config)

    def do_show_config(self, line):
        """
        Print the current configuration information.
        """
        print('Current configuration')
        print('  Public Key: {0}'.format(self.config['public']))
        print('  ZKM Server: {0}'.format(self.config['server']))
        print('  Last Check: {0}'.format(self.config['since']))
        print()

    def do_show_contacts(self, line):
        """
        Print the current list of contacts.
        """
        print('Contacts')
        for contact in self.contacts:
            print('  {0}: {1}'.format(contact, self.contacts[contact]))

        print()

    def do_create_message(self, line):
        """
        Create a new encrypted message using the public key associated with
        name.
        """
        line = line.split(' ')
        username = line[0]
        message = ' '.join(line[1:])
        their_public = self.contacts.get(username, None)

        if their_public is None:
            print('[-] No public key available for {0}.'.format(their_public))

        else:
            enc_msg = self._encrypt(self.config['secret'],
                                    self.config['public'],
                                    their_public,
                                    'message: {0}'.format(message))

            resp = send('POST', '/message', {'message': enc_msg})
            print('[+] {0}'.format(resp))

    def do_read_messages(self):
        """
        Get all messages and attempt to decrypt them.

        Use the since value stored in the configuration file. The server will
        return no more than the last 200 messages by default. This value is
        adjustable in the db.py script.
        """
        since = self.config.get('since', 1)

        resp = send('GET', '/messages', {'since': since})

        for enc_msg in resp:
            since = enc_msg['id']
            their_public, dec_msg = decrypt(enc_msg['message'])

            # Decryption was successful print the message
            if dec_msg.startswith('message: '):
                print_msg(their_public, dec_msg)

        # Update since value in the config
        self.config['since'] = since

    def do_EOF(self, line):
        """
        Usage quit | exit | ctrl-d

        Quit the Ada client.
        """
        return True

    def do_quit(self, line):
        return True

    def do_exit(self, line):
        return True

#-----------------------------------------------------------------------------
# Main Program
#-----------------------------------------------------------------------------
try:
    ZKMClient().cmdloop()

except Exception as e:
    print('[-] Error executing ZKM: {0}'.format(e))
<commit_msg>Made change from original name to ZKM.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2015 LCI Technology Group, LLC
# All rights reserved

import requests
import pysodium
import os
import base64
import json
import cmd

# CONSTANTS
HOMEDIR = os.path.expanduser("~")
ZKMDIR = os.path.join(HOMEDIR, '.zkm')
CONFIG = os.path.join(ZKMDIR, 'config')
CONTACT = os.path.join(ZKMDIR, 'contacts')

def load_json_file(filename):
    """
    Load a json file into a dictionary and return it.
    """
    with open(filename, 'rb') as f:
        data = f.read()

    try:
        return json.loads(data)
    except:
        return {}

def save_json_data(filename, data):
    """
    Save the given data to a file in JSON format.
    """
    with open(filename, 'wb') as f:
        f.write('{0}\n'.format(json.dumps(data)))

def encrypt(our_secret, our_public, their_public, msg):
    """
    Encrypt a message using the provided information.
    """
    their_public = base64.b64decode(their_public)
    our_secret = base64.b64decode(our_secret)
    nonce = pysodium.randombytes(pysodium.crypto_box_NONCEBYTES)
    enc = pysodium.crypto_box(msg, nonce, their_public, our_secret)

    nonce = base64.b64encode(nonce)
    enc = base64.b64encode(enc)

    # Return our public_key, nonce, and the encrypted message
    return ':'.join(our_public, nonce, enc)

def decrypt(our_secret, msg):
    """
    Decrypt a message using the provided information.
    """
    our_secret = base64.b64decode(our_secret)
    their_public, nonce, enc_msg = msg.split(':')
    their_public = base64.b64decode(their_public)
    nonce = base64.b64decode(nonce)
    enc_msg = base64.b64decode(enc_msg)
    dec_msg = pysodium.crypto_box_open(enc_msg, nonce, their_public, our_secret)

    # Return the sender's public key and the decrypted message.
    return their_public, dec_msg

def print_msg(contacts, their_public, msg):
    """
    Print the sender and message.

    Lookup the public key of the sender to see if they are in our
    contacts. If they are print the username, if not print the public key.
    """
    for username, contact_public in contacts.iteritems():
        if their_public == contact_public:
            sender = username
        else:
            sender = their_public

    print('{0}: {1}'.format(sender, msg))

def send(server, method, endpoint, data=None):
    """
    Send a message to the server and process the response.
    """
    url = '{0}{1}'.format(server, endpoint)

    if method == 'POST':
        resp = requests.post(url, data=data)

    if method == 'GET':
        resp = requests.get(url, params=data)

    if resp.status_code == 200:
        j = resp.json()
        if j['error'] is not None:
            print('[-] {0}'.format(j['error']))
            return None
        else:
            return j['response']
    else:
        print('[-] Server error: {0}'.format(resp.status_code))

def initialize():
    """
    Create a ~/.zkm directory with a config file inside.

    The config file will hold our public key, secret key, and since value.
    """
    if os.path.exists(ZKMDIR) is False:
        print('[+] Creating ZKM configuration directory.')
        os.mkdir(ZKMDIR, 0o750)

        print('[+] Creating new keypair.')
        our_public, our_secret = pysodium.crypto_sign_keypair()

        print('[+] Creating configuration file.')
        config = {'public': base64.b64encode(our_public),
                  'secret': base64.b64encode(our_secret),
                  'since': 1}

        save_json_data(CONFIG, config)
        os.chmod(CONFIG, 0o600)

        print('[+] Creating contacts file.')
        save_json_data(CONTACT, {})
        os.chmod(CONTACT, 0o600)

    else:
        print('[-] ZKM configuration directory already exists.')

class ZKMClient(cmd.Cmd):
    """
    ZKM: A zero knowledge messaging system.
    """
    prompt = 'zkm> '

    # Functions used in the interactive command prompt.
    def preloop(self):
        """
        Initialize the ZKM client if necessary.
        """
        try:
            self.config = load_json_file(CONFIG)
        except:
            print('[-] ZKM not initialized yet.')
            initialize()

        try:
            self.contacts = load_json_file(CONTACT)
        except Exception:
            print('[-] Could not load contacts file.')
            self.contacts = []

    def postloop(self):
        save_json_data(CONFIG, self.config)
        save_json_data(CONTACT, self.contacts)

    def do_add_contact(self, line):
        """
        Add a new contact to the contact list.
        """
        name, their_public = line.split(' ')
        self.contacts[name] = their_public
        save_json_data(CONTACT, self.contacts)

    def do_del_contact(self, name):
        """
        Remove a contact from the contact list.
        """
        self.contacts.pop(name, None)
        save_json_data(CONTACT, self.contacts)

    def do_connect(self, line):
        """
        Define the server we want to connect to for messages.
        """
        self.config['server'] = line
        save_json_data(CONFIG, self.config)

    def do_show_config(self, line):
        """
        Print the current configuration information.
        """
        print('Current configuration')
        print('  Public Key: {0}'.format(self.config['public']))
        print('  ZKM Server: {0}'.format(self.config['server']))
        print('  Last Check: {0}'.format(self.config['since']))
        print()

    def do_show_contacts(self, line):
        """
        Print the current list of contacts.
        """
        print('Contacts')
        for contact in self.contacts:
            print('  {0}: {1}'.format(contact, self.contacts[contact]))

        print()

    def do_create_message(self, line):
        """
        Create a new encrypted message using the public key associated with
        name.
        """
        line = line.split(' ')
        username = line[0]
        message = ' '.join(line[1:])
        their_public = self.contacts.get(username, None)

        if their_public is None:
            print('[-] No public key available for {0}.'.format(their_public))

        else:
            enc_msg = self._encrypt(self.config['secret'],
                                    self.config['public'],
                                    their_public,
                                    'message: {0}'.format(message))

            resp = send('POST', '/message', {'message': enc_msg})
            print('[+] {0}'.format(resp))

    def do_read_messages(self):
        """
        Get all messages and attempt to decrypt them.

        Use the since value stored in the configuration file. The server will
        return no more than the last 200 messages by default. This value is
        adjustable in the db.py script.
        """
        since = self.config.get('since', 1)

        resp = send('GET', '/messages', {'since': since})

        for enc_msg in resp:
            since = enc_msg['id']
            their_public, dec_msg = decrypt(enc_msg['message'])

            # Decryption was successful print the message
            if dec_msg.startswith('message: '):
                print_msg(their_public, dec_msg)

        # Update since value in the config
        self.config['since'] = since

    def do_EOF(self, line):
        """
        Usage quit | exit | ctrl-d

        Quit the ZKM client.
        """
        return True

    def do_quit(self, line):
        return True

    def do_exit(self, line):
        return True

#-----------------------------------------------------------------------------
# Main Program
#-----------------------------------------------------------------------------
try:
    ZKMClient().cmdloop()

except Exception as e:
    print('[-] Error executing ZKM: {0}'.format(e))
<|endoftext|>


<commit_before>import logging
import uuid
from bisect import insort
from datetime import datetime, timedelta
from pytz import utc
import parsedatetime

from errbot import BotPlugin, botcmd
from errbot.version import VERSION

__author__ = 'kdknowlton'

DEFAULT_POLL_INTERVAL = 60 * 5 # Five minutes

class Reminder(object):
    def __init__(self, date, message, target, is_user):
        self.id = uuid.uuid4().hex
        if not date.tzinfo:
            date = utc.localize(date)
        self.date = date
        self.message = message
        self.target = target
        self.is_user = is_user
        self.sent = False

    def __lt__(self, other):
        if hasattr(other, 'date'):
            return self.date < other.date
        else:
            raise NotImplemented

    def __gt__(self, other):
        if hasattr(other, 'date'):
            return self.date > other.date
        else:
            raise NotImplemented

    def send(self, bot):
        logging.info(
            'Sending reminder to {target}:\n{message}'.format(
                target=self.target,
                message=self.message
            )
        )
        message_type = 'chat' if self.is_user else 'groupchat'
        bot.send(
            self.target,
            "Hi! You asked me to remind you to {message}".format(message=self.message),
            message_type=message_type
        )
        self.sent = True
        bot.store_reminder(self)

    def is_ready(self, date=None):
        if date is None:
            date = datetime.now(utc)
        return date > self.date and not self.sent

class ReminderPlugin(BotPlugin):
    min_err_version = '1.6.0'
    cal = parsedatetime.Calendar()

    def get_configuration_template(self):
        return {'POLL_INTERVAL': DEFAULT_POLL_INTERVAL}

    def configure(self, configuration):
        if configuration:
            if type(configuration) != dict:
                raise Exception('Wrong configuration type')

            if not configuration.has_key('POLL_INTERVAL'):
                raise Exception('Wrong configuration type, it should contain POLL_INTERVAL')
            if len(configuration) > 1:
                raise Exception('What else did you try to insert in my config ?')
            try:
                int(configuration['POLL_INTERVAL'])
            except:
                raise Exception('POLL_INTERVAL must be an integer')
        super(ReminderPlugin, self).configure(configuration)

    def activate(self):
        super(ReminderPlugin, self).activate()
        self.send_reminders()
        self.start_poller(
            self.config['POLL_INTERVAL'] if self.config else DEFAULT_POLL_INTERVAL,
            self.send_reminders
        )

    def store_reminder(self, reminder):
        all_reminders = self.get('all_reminders', {})
        all_reminders[reminder.id] = reminder
        self['all_reminders'] = all_reminders

    def add_reminder(self, date, message, target, is_user=True):
        new = Reminder(date, message, target, is_user)
        self.store_reminder(new)
        return new

    # def get_reminders(self, target, is_user=True):
    #     if is_user:
    #         reminder_dict = self.get('user_reminders', {})
    #     else:
    #         reminder_dict = self.get('chatroom_reminders', {})
    #     return reminder_dict.get(target, [])

    def get_all_reminders(self):
        return self.get('all_reminders', {}).values()

    def send_reminders(self):
        logging.info('Sending all reminders.')
        for reminder in self.get_all_reminders():
            if reminder.is_ready():
                reminder.send(self)

    @botcmd
    def reminders_echo(self, mess, args):
        return args

    @botcmd
    def time(self, mess, args):
        return "The current local date and time is: " + datetime.now().strftime("%a, %b %d, %Y %H:%M:%S")

    @botcmd(split_args_with=' ')
    def remind_me(self, mess, args):
        if "to" not in args:
            return "Sorry, I didn't understand that."

        date_end = args.index('to')
        date_list = args[:date_end]
        date_string = " ".join(date_list)
        date = datetime(*(ReminderPlugin.cal.parse(date_string)[0])[:6])
        message = " ".join(args[date_end + 1:])
        is_user = mess.getType() == 'chat'
        target = mess.getFrom().getStripped()
        self.add_reminder(date, message, target, is_user)
        return "Ok, I'll remind you to {message} at {date}.".format(message=message, date=date)

    @botcmd(admin_only=True)
    def clear_all_reminders(self, mess, args):
        """WARNING: This will clear all reminders for all users and rooms!
        """
        self['user_reminders'] = {}
        self['chatroom_reminders'] = {}
        self['all_reminders'] = {}
        return 'All reminders have been cleared.'<commit_msg>Let the user know if we can't parse the date<commit_after>import logging
import uuid
from bisect import insort
from datetime import datetime, timedelta
from pytz import utc
import parsedatetime

from errbot import BotPlugin, botcmd
from errbot.version import VERSION

__author__ = 'kdknowlton'

DEFAULT_POLL_INTERVAL = 60 * 1 # Five minutes

class Reminder(object):
    def __init__(self, date, message, target, is_user):
        self.id = uuid.uuid4().hex
        if not date.tzinfo:
            date = utc.localize(date)
        self.date = date
        self.message = message
        self.target = target
        self.is_user = is_user
        self.sent = False

    def __lt__(self, other):
        if hasattr(other, 'date'):
            return self.date < other.date
        else:
            raise NotImplemented

    def __gt__(self, other):
        if hasattr(other, 'date'):
            return self.date > other.date
        else:
            raise NotImplemented

    def send(self, bot):
        logging.info(
            'Sending reminder to {target}:\n{message}'.format(
                target=self.target,
                message=self.message
            )
        )
        message_type = 'chat' if self.is_user else 'groupchat'
        bot.send(
            self.target,
            "Hi! You asked me to remind you to {message}".format(message=self.message),
            message_type=message_type
        )
        self.sent = True
        bot.store_reminder(self)

    def is_ready(self, date=None):
        if date is None:
            date = datetime.now(utc)
        return date > self.date and not self.sent

class ReminderPlugin(BotPlugin):
    min_err_version = '1.6.0'
    cal = parsedatetime.Calendar()

    def get_configuration_template(self):
        return {'POLL_INTERVAL': DEFAULT_POLL_INTERVAL}

    def configure(self, configuration):
        if configuration:
            if type(configuration) != dict:
                raise Exception('Wrong configuration type')

            if not configuration.has_key('POLL_INTERVAL'):
                raise Exception('Wrong configuration type, it should contain POLL_INTERVAL')
            if len(configuration) > 1:
                raise Exception('What else did you try to insert in my config ?')
            try:
                int(configuration['POLL_INTERVAL'])
            except:
                raise Exception('POLL_INTERVAL must be an integer')
        super(ReminderPlugin, self).configure(configuration)

    def activate(self):
        super(ReminderPlugin, self).activate()
        self.send_reminders()
        self.start_poller(
            self.config['POLL_INTERVAL'] if self.config else DEFAULT_POLL_INTERVAL,
            self.send_reminders
        )

    def store_reminder(self, reminder):
        all_reminders = self.get('all_reminders', {})
        all_reminders[reminder.id] = reminder
        self['all_reminders'] = all_reminders

    def add_reminder(self, date, message, target, is_user=True):
        new = Reminder(date, message, target, is_user)
        self.store_reminder(new)
        return new

    # def get_reminders(self, target, is_user=True):
    #     if is_user:
    #         reminder_dict = self.get('user_reminders', {})
    #     else:
    #         reminder_dict = self.get('chatroom_reminders', {})
    #     return reminder_dict.get(target, [])

    def get_all_reminders(self):
        return self.get('all_reminders', {}).values()

    def send_reminders(self):
        logging.info('Sending all reminders.')
        for reminder in self.get_all_reminders():
            if reminder.is_ready():
                reminder.send(self)

    @botcmd
    def reminders_echo(self, mess, args):
        return args

    @botcmd
    def time(self, mess, args):
        return "The current local date and time is: " + datetime.now().strftime("%a, %b %d, %Y %H:%M:%S")

    @botcmd(split_args_with=' ')
    def remind_me(self, mess, args):
        """Takes a message of the form '!remind me [when] to [what]' and stores the reminder."""
        if "to" not in args:
            return "Sorry, I didn't understand that."

        date_end = args.index('to')
        date_list = args[:date_end]
        date_string = " ".join(date_list)
        date_struct = ReminderPlugin.cal.parse(date_string)
        if ( date_struct[1] != 0 ):
            date = datetime(*(date_struct[0])[:6])
            message = " ".join(args[date_end + 1:])
            is_user = mess.getType() == 'chat'
            target = mess.getFrom().getStripped()
            self.add_reminder(date, message, target, is_user)
            return "Ok, I'll remind you to {message} at {date}.".format(message=message, date=date)
        else:
            return "I'm not sure when {date} is".format(date=date_string)

    @botcmd(admin_only=True)
    def clear_all_reminders(self, mess, args):
        """WARNING: This will clear all reminders for all users and rooms!
        """
        self['user_reminders'] = {}
        self['chatroom_reminders'] = {}
        self['all_reminders'] = {}
        return 'All reminders have been cleared.'<|endoftext|>


<commit_before>import json
from random import choice

class CobwebTree:

    # static variable for hashing concepts
    counter = 0
    root = None

    def __init__(self, tree=None):
        """
        The constructor.
        """
        if self.__class__.root == None:
            self.__class__.root = self
        self.concept_name = "Concept" + self._gensym()
        self.count = 0
        self.av_counts = {}
        self.children = []

        # check if the constructor is being used as a copy constructor
        if tree:
            self._update_counts_from_node(tree)

            for child in tree.children:
                self.children.append(self.__class__(child))

    def __hash__(self):
        return hash(self.concept_name)

    def _gensym(self):
        self.__class__.counter += 1
        return str(self.__class__.counter)

    def _shallow_copy(self):
        """
        Creates a copy of the current node and its children (but not their
        children)
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)

        for child in self.children:
            temp_child = self.__class__()
            temp_child._update_counts_from_node(child)
            temp.children.append(temp_child)

        return temp

    def _increment_counts(self, instance):
        """
        Increment the counts at the current node according to the specified
        instance.

        input:
            instance: {a1: v1, a2: v2, ...} - a hashtable of attr and values.
        """
        #TODO do counts need to be floats?
        self.count += 1
        for attr in instance:
            self.av_counts[attr] = self.av_counts.setdefault(attr,{})
            self.av_counts[attr][instance[attr]] = (self.av_counts[attr].get(
                instance[attr], 0) + 1)

    def _update_counts_from_node(self, node):
        """
        Increments the counts of the current node by the amount in the specified
        node.
        """
        self.count += node.count
        for attr in node.av_counts:
            for val in node.av_counts[attr]:
                self.av_counts[attr] = self.av_counts.setdefault(attr,{})
                self.av_counts[attr][val] = (self.av_counts[attr].get(val,0) +
                                     node.av_counts[attr][val])

    def _two_best_children(self, instance):
        """
        Returns the two best children to incorporate the instance
        into in terms of category utility.

        input:
            instance: {a1: v1, a2: v2,...} - a hashtable of attr. and values.
        output:
            (0.2,2),(0.1,3) - the category utility and indices for the two best
            children (the second tuple will be None if there is only 1 child).
        """
        if len(self.children) == 0:
            raise Exception("No children!")

        children_cu = [(self._cu_for_insert(child, instance), child) for child
                       in self.children]
        children_cu.sort(key=lambda x: x[0],reverse=True)

        if len(children_cu) == 1:
            return children_cu[0], None

        return children_cu[0], children_cu[1]

    def _cu_for_insert(self, child, instance):
        """
        Computer the category utility of adding the instance to the specified
        child.
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)
        temp._increment_counts(instance)

        for c in self.children:
            temp_child = self.__class__()
            temp_child._update_counts_from_node(c)
            temp.children.append(temp_child)
            if c == child:
                temp_child._increment_counts(instance)
        return temp._category_utility()

    def _create_new_child(self, instance):
        """
        Creates a new child (to the current node) with the counts initialized by
        the given instance.
        """
        new_child = self.__class__()
        new_child._increment_counts(instance)
        self.children.append(new_child)
        return new_child

    def _create_child_with_current_counts(self):
        """
        Creates a new child (to the current node) with the counts initialized by
        the current node's counts.
        """
        if self.count > 0:
            self.children.append(self.__class__(self))

    def _cu_for_new_child(self, instance):
        """
        Returns the category utility for creating a new child using the
        particular instance.
        """
        temp = self._shallow_copy()
        temp._increment_counts(instance)
        temp._create_new_child(instance)
        return temp._category_utility()

    def _merge(self, best1, best2, instance):
        """
        Merge the two specified nodes.

        input:
            best1: the best child
            best2: the second best child
        output:
            The new child formed from the merge
        """
        new_child = self.__class__()
        new_child._update_counts_from_node(best1)
        new_child._update_counts_from_node(best2)
        new_child.children.append(best1)
        new_child.children.append(best2)
        self.children.remove(best1)
        self.children.remove(best2)
        self.children.append(new_child)

        return new_child

    def _cu_for_merge(self, best1, best2, instance):
        """
        Returns the category utility for merging the two best children.
        NOTE! - I decided that testing a merge does not incorporate the latest
        instance, but waits for a second call of cobweb on the root. The
        original paper says that I should incorporate the instance into the
        merged node, but since we don't do something like this for split I
        didn't do it here. This gives the option to merge multiple nodes before
        incorporating the instance.

        input:
            best1: the best child in the children array.
            best2: the second best child in the children array.
        output:
            0.02 - the category utility for the merge of best1 and best2.
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)
        temp._increment_counts(instance)

        new_child = self.__class__()
        new_child._update_counts_from_node(best1)
        new_child._update_counts_from_node(best2)
        new_child._increment_counts(instance)
        temp.children.append(new_child)

        for c in self.children:
            if c == best1 or c == best2:
                continue
            temp_child = self.__class__()
            temp_child._update_counts_from_node(c)
            temp.children.append(temp_child)

        return temp._category_utility()

    def _split(self, best):
        """
        Split the best node and promote its children
        """
        self.children.remove(best)
        for child in best.children:
            self.children.append(child)

    def _cu_for_split(self, best):
        """
        Return the category utility for splitting the best child.

        input:
            best1: a child in the children array.
        output:
            0.03 - the category utility for the split of best1.
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)

        for c in self.children + best.children:
            if c == best:
                continue
            temp_child = self.__class__()
            temp_child._update_counts_from_node(c)
            temp.children.append(temp_child)

        return temp._category_utility()

    def _verify_counts(self):
        """
        Checks the property that the counts of the children sum to the same
        count as the parent. This is/was useful when debugging.
        """
        if len(self.children) == 0:
            return

        temp = {}
        temp_count = self.count
        for attr in self.av_counts:
            if attr not in temp:
                temp[attr] = {}
            for val in self.av_counts[attr]:
                temp[attr][val] = self.av_counts[attr][val]

        for child in self.children:
            temp_count -= child.count
            for attr in child.av_counts:
                assert attr in temp
                for val in child.av_counts[attr]:
                    assert val in temp[attr]
                    temp[attr][val] -= child.av_counts[attr][val]

        assert temp_count == 0

        for attr in temp:
            for val in temp[attr]:
                assert temp[attr][val] == 0.0

        for child in self.children:
            child._verify_counts()

    def _is_concept(self, instance):
        """
        Checks to see if the current node perfectly represents the instance (all
        of the attribute values the instance has are probability 1.0 and here
        are no extra attribute values).
        """
        for attribute in self.av_counts:
            for value in self.av_counts[attribute]:
                if (self.av_counts[attribute][value] / self.count) != 1.0:
                    return False
                if attribute not in instance:
                    return False
                if instance[attribute] != value:
                    return False

        for attribute in instance:
            if attribute not in self.av_counts:
                return False
            if instance[attribute] not in self.av_counts[attribute]:
                return False
            if ((self.av_counts[attribute][instance[attribute]] / self.count) !=
                1.0):
                return False

        return True

    def _match(self, instance):
        """
        Just a stub here, but it is used by Labyrinth later to match
        the concept.
        """
        return instance

    def _cobweb(self, instance):
        """
        Incrementally integrates an instance into the categorization tree
        defined by the current node. This function operates recursively to
        integrate this instance and uses category utility as the heuristic to
        make decisions.
        """
        if not self.children and self._is_concept(instance):
            self._increment_counts(instance)
            return self

        elif not self.children:
            self._create_child_with_current_counts()
            self._increment_counts(instance)
            return self._create_new_child(instance)

        else:
            #TODO is there a cleaner way to do this?
            best1, best2 = self._two_best_children(instance)
            action_cu, best_action = self._get_best_operation(instance, best1,
                                                              best2)
            best1_cu, best1 = best1
            if best2:
                best2_cu, best2 = best2

            if action_cu == 0.0 or best_action == 'best':
                self._increment_counts(instance)
                return best1._cobweb(instance)
            elif best_action == 'new':
                self._increment_counts(instance)
                return self._create_new_child(instance)
            elif best_action == 'merge':
                self._increment_counts(instance)
                new_child = self._merge(best1, best2, instance)
                return new_child._cobweb(instance)
            elif best_action == 'split':
                self._split(best1)
                return self._cobweb(instance)
            else:
                raise Exception("Should never get here.")

    def _get_best_operation(self, instance, best1, best2,
                            possible_ops=["best", "new", "merge", "split"]):
        """
        Given a set of possible operations, find the best and return its cu and
        the action name.
        """
        best1_cu, best1 = best1
        if best2:
            best2_cu, best2 = best2
        operations = []

        if "best" in possible_ops:
            operations.append((best1_cu,"best"))
        if "new" in possible_ops:
            operations.append((self._cu_for_new_child(instance),'new'))
        if "merge" in possible_ops and best2:
            operations.append((self._cu_for_merge(best1, best2,
                                                 instance),'merge'))
        if "split" in possible_ops and len(best1.children) > 0:
            operations.append((self._cu_for_split(best1),'split'))

        # pick the best operation
        operations.sort(reverse=True)

        return operations[0]

    def _cobweb_categorize(self, instance):
        """
        Sorts an instance in the categorization tree defined at the current
        node without modifying the counts of the tree.
        """
        if not self.children:
            return self

        best1, best2 = self._two_best_children(instance)
        action_cu, best_action = self._get_best_operation(instance, best1,
                                                          best2, ["best",
                                                                  "new"])
        best1_cu, best1 = best1

        if best_action == "new":
            return self
        elif best_action == "best":
            return best1._cobweb_categorize(instance)

    def _expected_correct_guesses(self):
        correct_guesses = 0.0
        for attr in self.av_counts:
            for val in self.av_counts[attr]:
                prob = ((1.0 * self.av_counts[attr][val]) / self.count)
                correct_guesses += (prob * prob)
        return correct_guesses

    def _category_utility(self):
        if len(self.children) == 0:
            return 0.0

        category_utility = 0.0

        for child in self.children:
            p_of_child = child.count / self.count
            category_utility += (p_of_child *
                                 (child._expected_correct_guesses()
                                  - self._expected_correct_guesses()))
        return category_utility / (1.0 * len(self.children))

    def _category_utility_old(self):
        """
        The category utility is a local heuristic calculation to determine if
        the split of instances across the children increases the ability to
        guess from the parent node.
        """
        if len(self.children) == 0:
            return 0.0

        category_utility = 0.0

        exp_parent_guesses = self._expected_correct_guesses()

        for child in self.children:
            p_of_child = child.count / self.count
            exp_child_guesses = child._expected_correct_guesses()
            category_utility += p_of_child * (exp_child_guesses -
                                              exp_parent_guesses)

        return category_utility / (1.0 * len(self.children))

    #def _expected_correct_guesses(self):
    #    """
    #    The number of attribute value guesses we would be expected to get
    #    correct using the current concept.
    #    """
    #    if self.count == 0:
    #        return 0.0

    #    exp_count = 0.0
    #    for attr in self.av_counts:
    #        for val in self.av_counts[attr]:
    #            if not isinstance(val, float):
    #                exp_count += (self.av_counts[attr][val] / self.count)**2
    #    return exp_count

    def _num_concepts(self):
        """
        Return the number of concepts contained in the tree defined by the
        current node.
        """
        children_count = 0
        for c in self.children:
           children_count += c._num_concepts()
        return 1 + children_count

    def _pretty_print(self, depth=0):
        """
        Prints the categorization tree.
        """
        ret = str(('\t' * depth) + "|-" + str(self.av_counts) + ":" +
                  str(self.count) + '\n')

        for c in self.children:
            ret += c._pretty_print(depth+1)

        return ret

    def __str__(self):
        """
        Converts the categorization tree into a string for printing"
        """
        return self._pretty_print()

    def ifit(self, instance):
        """
        Given an instance incrementally update the categorization tree.
        """
        self._cobweb(instance)

    def fit(self, list_of_instances):
        """
        Call incremental fit on each element in a list of instances.
        """
        for instance in list_of_instances:
            self.ifit(instance)

    def predict(self, instance):
        """
        Given an instance predict any missing attribute values without
        modifying the tree.
        """
        prediction = {}

        # make a copy of the instance
        for attr in instance:
            prediction[attr] = instance[attr]

        concept = self._cobweb_categorize(instance)

        for attr in concept.av_counts:
            if attr in prediction:
                continue

            values = []
            for val in concept.av_counts[attr]:
                values += [val] * concept.av_counts[attr][val]

            prediction[attr] = choice(values)

        return prediction

    def predict_all(self, instances):
        """
        Predicts missing attribute values of all instances in the give
        list.
        """
        predictions = [self.predict(instance) for instance in instances]
        return predictions

    def train_from_json(self, filename):
        json_data = open(filename, "r")
        instances = json.load(json_data)
        self.fit(instances)

        json_data.close()

if __name__ == "__main__":
    t = CobwebTree()
    t.train_from_json("cobweb_test.json")
    t._verify_counts()
    print(t)
    print()

    test = {}
    test['HeartChamber'] = "four"
    print(t.predict(test))

<commit_msg>removed some unnecessary code<commit_after>import json
from random import choice

class CobwebTree:

    # static variable for hashing concepts
    counter = 0
    root = None

    def __init__(self, tree=None):
        """
        The constructor.
        """
        if self.__class__.root == None:
            self.__class__.root = self
        self.concept_name = "Concept" + self._gensym()
        self.count = 0
        self.av_counts = {}
        self.children = []

        # check if the constructor is being used as a copy constructor
        if tree:
            self._update_counts_from_node(tree)

            for child in tree.children:
                self.children.append(self.__class__(child))

    def __hash__(self):
        return hash(self.concept_name)

    def _gensym(self):
        self.__class__.counter += 1
        return str(self.__class__.counter)

    def _shallow_copy(self):
        """
        Creates a copy of the current node and its children (but not their
        children)
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)

        for child in self.children:
            temp_child = self.__class__()
            temp_child._update_counts_from_node(child)
            temp.children.append(temp_child)

        return temp

    def _increment_counts(self, instance):
        """
        Increment the counts at the current node according to the specified
        instance.

        input:
            instance: {a1: v1, a2: v2, ...} - a hashtable of attr and values.
        """
        #TODO do counts need to be floats?
        self.count += 1
        for attr in instance:
            self.av_counts[attr] = self.av_counts.setdefault(attr,{})
            self.av_counts[attr][instance[attr]] = (self.av_counts[attr].get(
                instance[attr], 0) + 1)

    def _update_counts_from_node(self, node):
        """
        Increments the counts of the current node by the amount in the specified
        node.
        """
        self.count += node.count
        for attr in node.av_counts:
            for val in node.av_counts[attr]:
                self.av_counts[attr] = self.av_counts.setdefault(attr,{})
                self.av_counts[attr][val] = (self.av_counts[attr].get(val,0) +
                                     node.av_counts[attr][val])

    def _two_best_children(self, instance):
        """
        Returns the two best children to incorporate the instance
        into in terms of category utility.

        input:
            instance: {a1: v1, a2: v2,...} - a hashtable of attr. and values.
        output:
            (0.2,2),(0.1,3) - the category utility and indices for the two best
            children (the second tuple will be None if there is only 1 child).
        """
        if len(self.children) == 0:
            raise Exception("No children!")

        children_cu = [(self._cu_for_insert(child, instance), child) for child
                       in self.children]
        children_cu.sort(key=lambda x: x[0],reverse=True)

        if len(children_cu) == 1:
            return children_cu[0], None

        return children_cu[0], children_cu[1]

    def _cu_for_insert(self, child, instance):
        """
        Computer the category utility of adding the instance to the specified
        child.
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)
        temp._increment_counts(instance)

        for c in self.children:
            temp_child = self.__class__()
            temp_child._update_counts_from_node(c)
            temp.children.append(temp_child)
            if c == child:
                temp_child._increment_counts(instance)
        return temp._category_utility()

    def _create_new_child(self, instance):
        """
        Creates a new child (to the current node) with the counts initialized by
        the given instance.
        """
        new_child = self.__class__()
        new_child._increment_counts(instance)
        self.children.append(new_child)
        return new_child

    def _create_child_with_current_counts(self):
        """
        Creates a new child (to the current node) with the counts initialized by
        the current node's counts.
        """
        if self.count > 0:
            self.children.append(self.__class__(self))

    def _cu_for_new_child(self, instance):
        """
        Returns the category utility for creating a new child using the
        particular instance.
        """
        temp = self._shallow_copy()
        temp._increment_counts(instance)
        temp._create_new_child(instance)
        return temp._category_utility()

    def _merge(self, best1, best2, instance):
        """
        Merge the two specified nodes.

        input:
            best1: the best child
            best2: the second best child
        output:
            The new child formed from the merge
        """
        new_child = self.__class__()
        new_child._update_counts_from_node(best1)
        new_child._update_counts_from_node(best2)
        new_child.children.append(best1)
        new_child.children.append(best2)
        self.children.remove(best1)
        self.children.remove(best2)
        self.children.append(new_child)

        return new_child

    def _cu_for_merge(self, best1, best2, instance):
        """
        Returns the category utility for merging the two best children.
        NOTE! - I decided that testing a merge does not incorporate the latest
        instance, but waits for a second call of cobweb on the root. The
        original paper says that I should incorporate the instance into the
        merged node, but since we don't do something like this for split I
        didn't do it here. This gives the option to merge multiple nodes before
        incorporating the instance.

        input:
            best1: the best child in the children array.
            best2: the second best child in the children array.
        output:
            0.02 - the category utility for the merge of best1 and best2.
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)
        temp._increment_counts(instance)

        new_child = self.__class__()
        new_child._update_counts_from_node(best1)
        new_child._update_counts_from_node(best2)
        new_child._increment_counts(instance)
        temp.children.append(new_child)

        for c in self.children:
            if c == best1 or c == best2:
                continue
            temp_child = self.__class__()
            temp_child._update_counts_from_node(c)
            temp.children.append(temp_child)

        return temp._category_utility()

    def _split(self, best):
        """
        Split the best node and promote its children
        """
        self.children.remove(best)
        for child in best.children:
            self.children.append(child)

    def _cu_for_split(self, best):
        """
        Return the category utility for splitting the best child.

        input:
            best1: a child in the children array.
        output:
            0.03 - the category utility for the split of best1.
        """
        temp = self.__class__()
        temp._update_counts_from_node(self)

        for c in self.children + best.children:
            if c == best:
                continue
            temp_child = self.__class__()
            temp_child._update_counts_from_node(c)
            temp.children.append(temp_child)

        return temp._category_utility()

    def _verify_counts(self):
        """
        Checks the property that the counts of the children sum to the same
        count as the parent. This is/was useful when debugging.
        """
        if len(self.children) == 0:
            return

        temp = {}
        temp_count = self.count
        for attr in self.av_counts:
            if attr not in temp:
                temp[attr] = {}
            for val in self.av_counts[attr]:
                temp[attr][val] = self.av_counts[attr][val]

        for child in self.children:
            temp_count -= child.count
            for attr in child.av_counts:
                assert attr in temp
                for val in child.av_counts[attr]:
                    assert val in temp[attr]
                    temp[attr][val] -= child.av_counts[attr][val]

        assert temp_count == 0

        for attr in temp:
            for val in temp[attr]:
                assert temp[attr][val] == 0.0

        for child in self.children:
            child._verify_counts()

    def _is_concept(self, instance):
        """
        Checks to see if the current node perfectly represents the instance (all
        of the attribute values the instance has are probability 1.0 and here
        are no extra attribute values).
        """
        for attribute in self.av_counts:
            for value in self.av_counts[attribute]:
                if (self.av_counts[attribute][value] / self.count) != 1.0:
                    return False
                if attribute not in instance:
                    return False
                if instance[attribute] != value:
                    return False

        for attribute in instance:
            if attribute not in self.av_counts:
                return False
            if instance[attribute] not in self.av_counts[attribute]:
                return False
            if ((self.av_counts[attribute][instance[attribute]] / self.count) !=
                1.0):
                return False

        return True

    def _match(self, instance):
        """
        Just a stub here, but it is used by Labyrinth later to match
        the concept.
        """
        return instance

    def _cobweb(self, instance):
        """
        Incrementally integrates an instance into the categorization tree
        defined by the current node. This function operates recursively to
        integrate this instance and uses category utility as the heuristic to
        make decisions.
        """
        if not self.children and self._is_concept(instance):
            self._increment_counts(instance)
            return self

        elif not self.children:
            self._create_child_with_current_counts()
            self._increment_counts(instance)
            return self._create_new_child(instance)

        else:
            #TODO is there a cleaner way to do this?
            best1, best2 = self._two_best_children(instance)
            action_cu, best_action = self._get_best_operation(instance, best1,
                                                              best2)
            best1_cu, best1 = best1
            if best2:
                best2_cu, best2 = best2

            if action_cu == 0.0 or best_action == 'best':
                self._increment_counts(instance)
                return best1._cobweb(instance)
            elif best_action == 'new':
                self._increment_counts(instance)
                return self._create_new_child(instance)
            elif best_action == 'merge':
                self._increment_counts(instance)
                new_child = self._merge(best1, best2, instance)
                return new_child._cobweb(instance)
            elif best_action == 'split':
                self._split(best1)
                return self._cobweb(instance)
            else:
                raise Exception("Should never get here.")

    def _get_best_operation(self, instance, best1, best2,
                            possible_ops=["best", "new", "merge", "split"]):
        """
        Given a set of possible operations, find the best and return its cu and
        the action name.
        """
        best1_cu, best1 = best1
        if best2:
            best2_cu, best2 = best2
        operations = []

        if "best" in possible_ops:
            operations.append((best1_cu,"best"))
        if "new" in possible_ops:
            operations.append((self._cu_for_new_child(instance),'new'))
        if "merge" in possible_ops and best2:
            operations.append((self._cu_for_merge(best1, best2,
                                                 instance),'merge'))
        if "split" in possible_ops and len(best1.children) > 0:
            operations.append((self._cu_for_split(best1),'split'))

        # pick the best operation
        operations.sort(reverse=True)

        return operations[0]

    def _cobweb_categorize(self, instance):
        """
        Sorts an instance in the categorization tree defined at the current
        node without modifying the counts of the tree.
        """
        if not self.children:
            return self

        best1, best2 = self._two_best_children(instance)
        action_cu, best_action = self._get_best_operation(instance, best1,
                                                          best2, ["best",
                                                                  "new"])
        best1_cu, best1 = best1

        if best_action == "new":
            return self
        elif best_action == "best":
            return best1._cobweb_categorize(instance)

    def _expected_correct_guesses(self):
        correct_guesses = 0.0
        for attr in self.av_counts:
            for val in self.av_counts[attr]:
                prob = ((1.0 * self.av_counts[attr][val]) / self.count)
                correct_guesses += (prob * prob)
        return correct_guesses

    def _category_utility(self):
        if len(self.children) == 0:
            return 0.0

        category_utility = 0.0

        for child in self.children:
            p_of_child = child.count / self.count
            category_utility += (p_of_child *
                                 (child._expected_correct_guesses()
                                  - self._expected_correct_guesses()))
        return category_utility / (1.0 * len(self.children))

    def _num_concepts(self):
        """
        Return the number of concepts contained in the tree defined by the
        current node.
        """
        children_count = 0
        for c in self.children:
           children_count += c._num_concepts()
        return 1 + children_count

    def _pretty_print(self, depth=0):
        """
        Prints the categorization tree.
        """
        ret = str(('\t' * depth) + "|-" + str(self.av_counts) + ":" +
                  str(self.count) + '\n')

        for c in self.children:
            ret += c._pretty_print(depth+1)

        return ret

    def __str__(self):
        """
        Converts the categorization tree into a string for printing"
        """
        return self._pretty_print()

    def ifit(self, instance):
        """
        Given an instance incrementally update the categorization tree.
        """
        self._cobweb(instance)

    def fit(self, list_of_instances):
        """
        Call incremental fit on each element in a list of instances.
        """
        for instance in list_of_instances:
            self.ifit(instance)

    def predict(self, instance):
        """
        Given an instance predict any missing attribute values without
        modifying the tree.
        """
        prediction = {}

        # make a copy of the instance
        for attr in instance:
            prediction[attr] = instance[attr]

        concept = self._cobweb_categorize(instance)

        for attr in concept.av_counts:
            if attr in prediction:
                continue

            values = []
            for val in concept.av_counts[attr]:
                values += [val] * concept.av_counts[attr][val]

            prediction[attr] = choice(values)

        return prediction

    def predict_all(self, instances):
        """
        Predicts missing attribute values of all instances in the give
        list.
        """
        predictions = [self.predict(instance) for instance in instances]
        return predictions

    def train_from_json(self, filename):
        json_data = open(filename, "r")
        instances = json.load(json_data)
        self.fit(instances)

        json_data.close()

if __name__ == "__main__":
    t = CobwebTree()
    t.train_from_json("cobweb_test.json")
    t._verify_counts()
    print(t)
    print()

    test = {}
    test['HeartChamber'] = "four"
    print(t.predict(test))

<|endoftext|>


<commit_before>import importlib
import traceback
import modular
importlib.reload(modular)

ALIASES = {
        'home': 'run 6s w 2s e 2n\nopen w\nw',
        'rt vassendar': 'run 4s d w d 2w d 2n 2e\nopen s\ns\nopen d\nrun 5d\nopen w\nw\nrun 8n w 2s 6w\nopen w\nrun 11w 3n 3w\nopen w\nrun 5w\nrun 3n 5w',
        'rt wgate': 'run 2s 3w\nopen w\nw',
        'rt sehaire': 'run w u 6w 2n 3w s 6w s 6w 2n 5w 5n w n w n 4w n e',
        }

TRIGGERS = {
        'You feel a little cleaner, but are still very dirty.': 'bathe',
        'You feel a little cleaner.': 'bathe',
        'You feel a little cleaner; almost perfect.': 'bathe',
        'You are no longer hungry.': '!',
        'You are no longer thirsty.': '!',
        'You are starved, and near death.  EAT SOMETHING!': 'quit\ny',
        'You are dehydrated, and near death.  DRINK SOMETHING!': 'quit\ny',
        'YOU ARE DYING OF THIRST!': 'quit\ny',
        'YOU ARE DYING OF HUNGER!': 'quit\ny',
        }

class Coffee(modular.ModularClient):
    def __init__(self, mud, name):
        from modules.logging import Logging
        from modules.eval import Eval
        from modules.repeat import Repeat
        from modules.mapper import Mapper

        self.name = name
        self.logfname = '{}.log'.format(name)
        self.mapfname = '{}.map'.format(name)

        self.modules = {}
        mods = {
                'eval': (Eval, []),
                'repeat': (Repeat, []),
                'logging': (Logging, [self.logfname]),
                'mapper': (Mapper, [False, self.mapfname]),
                }
        if name == 'grumpy':
            from modules.gaoler import Gaoler
            mods['gaoler'] = (Gaoler, [])
        for modname, module in mods.items():
            try:
                constructor, args = module
                args = [mud] + args
                print("Constructing", constructor, "with", repr(args))
                self.modules[modname] = constructor(*args)
            except Exception:
                traceback.print_exc()

        super().__init__(mud)

        self.aliases.update(ALIASES)
        self.triggers.update(TRIGGERS)

        if name == 'cizra':
            self.triggers.update({
                    '(\w+): A closed door': lambda mud, matches: 'open ' + matches[0],
                    '(\w+) : A closed door': lambda mud, matches: 'open ' + matches[0],
                    'You point at .*, shooting forth a magic missile!': 'mm',
                    'You point at .*, shooting forth a blazing fireball!': 'ff',
                    'You are too close to .* to use Fireball.': 'mm',
                    'You point at .*, but nothing more happens.': 'ff',
                    'You point at .*, but fizzle the spell.': 'mm',
                    'You attempt to invoke magical protection, but fail.': 'cast mage armor',
                    })

    def getHostPort(self):
        return 'coffeemud.net', 2323

def getClass():
    return Coffee
<commit_msg>update Coffee<commit_after>import importlib
import traceback
import modular
importlib.reload(modular)

ALIASES = {
        'home': 'run 6s w 2s e 2n\nopen w\nw',
        'rt vassendar': 'run 4s d w d 2w d 2n 2e\nopen s\ns\nopen d\nrun 5d\nopen w\nw\nrun 8n w 2s 6w\nopen w\nrun 11w 3n 3w\nopen w\nrun 5w\nrun 3n 5w',
        'rt wgate': 'run 2s 3w\nopen w\nw',
        'rt sehaire': 'run w u 6w 2n 3w s 6w s 6w 2n 5w 5n w n w n 4w n e',
        }

TRIGGERS = {
        'You feel a little cleaner, but are still very dirty.': 'bathe',
        'You feel a little cleaner.': 'bathe',
        'You feel a little cleaner; almost perfect.': 'bathe',
        'You are no longer hungry.': '!',
        'You are no longer thirsty.': '!',
        'You are starved, and near death.  EAT SOMETHING!': 'quit\ny',
        'You are dehydrated, and near death.  DRINK SOMETHING!': 'quit\ny',
        'YOU ARE DYING OF THIRST!': 'quit\ny',
        'YOU ARE DYING OF HUNGER!': 'quit\ny',
        }

class Coffee(modular.ModularClient):
    def __init__(self, mud, name):
        from modules.logging import Logging
        from modules.eval import Eval
        from modules.repeat import Repeat
        from modules.mapper import Mapper

        self.name = name
        self.logfname = '{}.log'.format(name)
        self.mapfname = '{}.map'.format(name)

        self.modules = {}
        mods = {
                'eval': (Eval, []),
                'repeat': (Repeat, []),
                'logging': (Logging, [self.logfname]),
                'mapper': (Mapper, [False, self.mapfname]),
                }
        if name == 'grumpy':
            from modules.gaoler import Gaoler
            mods['gaoler'] = (Gaoler, [])
        for modname, module in mods.items():
            try:
                constructor, args = module
                args = [mud] + args
                print("Constructing", constructor, "with", repr(args))
                self.modules[modname] = constructor(*args)
            except Exception:
                traceback.print_exc()

        super().__init__(mud)

        self.aliases.update(ALIASES)
        self.triggers.update(TRIGGERS)

        if name == 'cizra':
            self.triggers.update({
                    '(\w+): A closed door': lambda mud, matches: 'open ' + matches[0],
                    '(\w+) : A closed door': lambda mud, matches: 'open ' + matches[0],
                    'You point at .*, shooting forth a magic missile!': 'mm',
                    'You point at .*, shooting forth a blazing fireball!': 'ff',
                    'You are too close to .* to use Fireball.': 'mm',
                    'You point at .*, but nothing more happens.': 'ff',
                    'You point at .*, but fizzle the spell.': 'mm',
                    'You attempt to invoke magical protection, but fail.': 'cast mage armor',
                    'You shout combatively, but nothing more happens.': 'cast combat precognition',
                    })

    def getHostPort(self):
        return 'coffeemud.net', 2323

def getClass():
    return Coffee
<|endoftext|>


<commit_before># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
from __future__ import absolute_import
import re
import json
import copy
import xml.etree.ElementTree as ET

from svtplay_dl.service import Service, OpenGraphThumbMixin
from svtplay_dl.utils.urllib import urljoin
from svtplay_dl.fetcher.rtmp import RTMP
from svtplay_dl.fetcher.hls import hlsparse
from svtplay_dl.log import log
from svtplay_dl.error import ServiceError
from svtplay_dl.subtitle import subtitle

class Urplay(Service, OpenGraphThumbMixin):
    supported_domains = ['urplay.se', 'ur.se']

    def __init__(self, url):
        Service.__init__(self, url)
        self.subtitle = None

    def get(self, options):
        data = self.get_urldata()
        match = re.search(r"urPlayer.init\((.*)\);", data)
        if not match:
            yield ServiceError("Can't find json info")
            return

        if self.exclude(options):
            yield ServiceError("Excluding video")
            return

        data = match.group(1)
        jsondata = json.loads(data)
        if len(jsondata["subtitles"]) > 0:
            yield subtitle(copy.copy(options), "tt", jsondata["subtitles"][0]["file"].split(",")[0])
        if "streamer" in jsondata["streaming_config"]:
            basedomain = jsondata["streaming_config"]["streamer"]["redirect"]
        else:
            lbjson = self.http.request("get", jsondata["streaming_config"]["loadbalancer"]).text
            lbjson = json.loads(lbjson)
            basedomain = lbjson["redirect"]
        http = "http://%s/%s" % (basedomain, jsondata["file_http"])
        hd = None
        if len(jsondata["file_http_hd"]) > 0:
            http_hd = "http://%s/%s" % (basedomain, jsondata["file_http_hd"])
            hls_hd = "%s%s" % (http_hd, jsondata["streaming_config"]["http_streaming"]["hls_file"])
            hd = True
        hls = "%s%s" % (http, jsondata["streaming_config"]["http_streaming"]["hls_file"])
        streams = hlsparse(options, self.http.request("get", hls), hls)
        for n in list(streams.keys()):
            yield streams[n]
        if hd:
            streams = hlsparse(options, self.http.request("get", hls_hd), hls_hd)
            for n in list(streams.keys()):
                yield streams[n]

    def scrape_episodes(self, options):
        res = []
        for relurl in re.findall(r'<a class="puff tv video"\s+title="[^"]*"\s+href="([^"]*)"',
                                 self.get_urldata()):
            res.append(urljoin(self.url, relurl.replace("&", "&")))

        if options.all_last != -1:
            res = res[-options.all_last:]

        return res

    def find_all_episodes(self, options):
        match = re.search(r'<link rel="alternate" type="application/rss\+xml" [^>]*href="([^"]+)"',
                          self.get_urldata())
        if match is None:
            log.info("Couldn't retrieve episode list as rss, trying to scrape")
            return self.scrape_episodes(options)

        url = "http://urplay.se%s" % match.group(1).replace("&", "&")
        xml = ET.XML(self.http.request("get", url).content)

        episodes = [x.text for x in xml.findall(".//item/link")]
        episodes_new = []
        n = 0
        for i in episodes:
            if n == options.all_last:
                break
            if i not in episodes_new:
                episodes_new.append(i)
            n += 1
        return episodes_new
<commit_msg>urplay: support for betaplay<commit_after># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
from __future__ import absolute_import
import re
import json
import copy
import xml.etree.ElementTree as ET

from svtplay_dl.service import Service, OpenGraphThumbMixin
from svtplay_dl.utils.urllib import urljoin
from svtplay_dl.fetcher.rtmp import RTMP
from svtplay_dl.fetcher.hls import hlsparse
from svtplay_dl.log import log
from svtplay_dl.error import ServiceError
from svtplay_dl.subtitle import subtitle

class Urplay(Service, OpenGraphThumbMixin):
    supported_domains = ['urplay.se', 'ur.se', 'betaplay.ur.se']

    def __init__(self, url):
        Service.__init__(self, url)
        self.subtitle = None

    def get(self, options):
        data = self.get_urldata()
        match = re.search(r"urPlayer.init\((.*)\);", data)
        if not match:
            yield ServiceError("Can't find json info")
            return

        if self.exclude(options):
            yield ServiceError("Excluding video")
            return

        data = match.group(1)
        jsondata = json.loads(data)
        if len(jsondata["subtitles"]) > 0:
            yield subtitle(copy.copy(options), "tt", jsondata["subtitles"][0]["file"].split(",")[0])
        if "streamer" in jsondata["streaming_config"]:
            basedomain = jsondata["streaming_config"]["streamer"]["redirect"]
        else:
            lbjson = self.http.request("get", jsondata["streaming_config"]["loadbalancer"]).text
            lbjson = json.loads(lbjson)
            basedomain = lbjson["redirect"]
        http = "http://%s/%s" % (basedomain, jsondata["file_http"])
        hd = None
        if len(jsondata["file_http_hd"]) > 0:
            http_hd = "http://%s/%s" % (basedomain, jsondata["file_http_hd"])
            hls_hd = "%s%s" % (http_hd, jsondata["streaming_config"]["http_streaming"]["hls_file"])
            hd = True
        hls = "%s%s" % (http, jsondata["streaming_config"]["http_streaming"]["hls_file"])
        streams = hlsparse(options, self.http.request("get", hls), hls)
        for n in list(streams.keys()):
            yield streams[n]
        if hd:
            streams = hlsparse(options, self.http.request("get", hls_hd), hls_hd)
            for n in list(streams.keys()):
                yield streams[n]

    def scrape_episodes(self, options):
        res = []
        for relurl in re.findall(r'<a class="puff tv video"\s+title="[^"]*"\s+href="([^"]*)"',
                                 self.get_urldata()):
            res.append(urljoin(self.url, relurl.replace("&", "&")))

        if options.all_last != -1:
            res = res[-options.all_last:]

        return res

    def find_all_episodes(self, options):
        match = re.search(r'<link rel="alternate" type="application/rss\+xml" [^>]*href="([^"]+)"',
                          self.get_urldata())
        if match is None:
            log.info("Couldn't retrieve episode list as rss, trying to scrape")
            return self.scrape_episodes(options)

        url = "http://urplay.se%s" % match.group(1).replace("&", "&")
        xml = ET.XML(self.http.request("get", url).content)

        episodes = [x.text for x in xml.findall(".//item/link")]
        episodes_new = []
        n = 0
        for i in episodes:
            if n == options.all_last:
                break
            if i not in episodes_new:
                episodes_new.append(i)
            n += 1
        return episodes_new
<|endoftext|>