<commit_before># -*- coding: utf-8 -*-
import sys, time, random
sys.path.append('/var/python/libs/')
sys.path.append('/var/python/')
sys.path.append('/home/python/libs/')
sys.path.append('/home/python/')

from flaskJSONRPCServer import flaskJSONRPCServer

class mySharedMethods:
   def random(self, mult=65536):
      # Sipmly return random value (0..mult)
      return int(random.random()*mult)

def echo(data='Hello world!'):
   # Simply echo
   return data

def myip(_connection=None):
   # Return client's IP
   return 'Hello, %s!'%(_connection.ip)

def setcookie(_connection=None):
   # Set cookie to client
   _connection.cookiesOut.append({'name':'myTestCookie', 'value':'Your IP is %s'%_connection.ip})
   return 'Setted'

def block():
   # Test for notification request. When it fully implemented, client must  not wait for compliting this function
   time.sleep(10)
   return 'ok'

def stats(_connection=None):
   #calculate connections per second
   tArr1={'connPerSec_now':_connection.server.connPerMinute.count/60.0, 'connPerSec_old':_connection.server.connPerMinute.oldCount/60.0, 'speedStats':{}}
   #calculate spped stats
   for k, v in _connection.server.speedStats.items():
      tArr1['speedStats'][k+'_min']=min(v)
      tArr1['speedStats'][k+'_max']=max(v)
      tArr1['speedStats'][k+'_average']=sum(v)/float(len(v))
   return tArr1

def big(_connection=None):
   _connection.allowCompress=True #allow compression for this method
   s="""
They say that if you sat an infinite number of monkeys with an infinite number of typewriters, then they would write the complete works of Shakespeare. This is known as the Infinite Monkey Theorem. I always feel it would be nice to try the actual experiment out. I mean all these monkeys with typewriters would be kinda cool.

However wouldn't it be easier to get one monkey and one typewriter and give him an infinite amount of time to see how he goes?

Although in theory it sounds simpler, getting hold of even a single monkey is not that easy.

But what if we created a simulated monkey in Python? Now we are talking! That would be easy, and unlike a real monkey our simulated one could work all day and all night. No food required.

Sounds like a plan for a blog post.

However I want to give our virtual monkey a chance. If he managed to type the complete works of Shakespeare without worrying about punctuation, we would be happy right? Same goes for capital letters, let's just ignore them for now. We want to give him a fighting chance of completing his task.

To help you get started I have taken a text file of the complete works of Shakespeare, and removed all punctuation and made all letters small. If any of you are thinking I did that by hand, think again. Python is your friend for automating tasks like that.

Because the complete works is a large tome, I have also included a similar file, with just Hamlet. You can choose which one you want your monkey to attempt!

As with all my blog posts I will give you the complete program first of all. Have a read through this and see how much of it you understand. Try and figure out those sections you don’t. I will then go through the program a line at a time explaining it in more detail.
We do this by using scriptRead.count(monkeyTyped) >= 1 What does that even mean? Well we take the complete works of Shakespeare stored in scriptRead, and we do a count of how many times the thing in the brackets appears. If we put monkeyTyped into the brackets then it will count how many times monkeyCount appeared in scriptRead. We only need it to appear once but if it appears more than once, then that is ok as well. We check this with the greater than or equal to symbol >=.

But the line also has monkeyTyped = ''. Why is that? Well when we start off this while loop the monkey has not typed anything, as we have set the variable to = '' , which means an empty string. This would mean our while loop would be false straight away, and our monkey would never start typing.

To get around this we us an or condition. We check that monkeyTyped is '' or scriptRead.count(monkeyTyped) >= 1

The next line is the line which does the actual typing.
We take monkeyTyped and we add the result of (random.choice(string.ascii_lowercase + ' ')) onto it and save the result as monkeyTyped.

What on earth is (random.choice(string.ascii_lowercase + ' ')) ?

Well we want to simulate the monkey typing on a keyboard. We have already said to make it easy we will not ask for punctuation, although we should ask him for a space between words. Thats not too much to ask is it? We also said we will not be too concerned about it being in capital letters.

So random.choice picks something from inside the brackets by random selection. So what have we put in the brackets. The first thing is

string.ascii_lowercase

This creates a string of the lower case ascii letters. i.e

abcdefghijklmnopqrstuvwxyz

We also add onto the end of that a space using ' '

Now remember this has a space between the two speech-marks. This is different from when we are creating a blank string, which doesn't have a space.

Once he has typed a word, the loop will start again and check if what has been typed is in the complete works. If it is he will continue. If not he will go to the else statement.
Here we return what he has typed. However remember we only get to this place if what is typed is not in the complete works. That's no use to us! We will need to remove the last letter of what he typed, as it was that last letter we know to have made perfect Shakespeare prose into gibberish.

So we return monkeyTyped[:-1] which returns everything up to the last letter but not the last letter.

Ok back to where we were in the main program.

Remember we said we would keep track of the number of key presses? Well lets update that now. We know the monkey has just typed a word, so it is easy to determine the size of that word and increase keyPresses by that amount. Oh but the monkey typed a letter which turned his prose into gibberish. We should add that on as well as we want our count to be accurate!
This line checks to see if the number of key presses he has made is a multiple of 10,000. The % checks if keyPresses / 10,000 has no remainder i.e. it is a multiple of 10,000. If it is then it will print the number of key-presses made. This only works if the key-presses are an exact multiple of 10,000. There will be times when the monkey is in the middle of a word when this happen, so it will not report back. However it reports enough to make you realise the monkey is still working and not fallen asleep.

Remember to press F5 to save and run your program.

That is the end of the program. All you need now it to download the complete works of Shakespeare or just Hamlet if thats what you want to use. Use the links below to do that.
   """
   return s

if __name__=='__main__':
   print 'Running api..'
   # Creating instance of server
   #    <blocking>      set is this server async
   #    <cors>          switch auto CORS support
   #    <gevent>        switch to using Gevent as backend
   #    <debug>         switch to logging connection's info from Flask
   #    <log>           switch to logging debug info from flaskJSONRPCServer
   #    <fallback>      switch auto fallback to JSONP on GET requests
   #    <allowCompress> switch auto compression
   server=flaskJSONRPCServer(("0.0.0.0", 7001), blocking=False, cors=True, gevent=False, debug=False, log=True, fallback=True, allowCompress=False)
   # Register dispatcher for all methods of instance
   server.registerInstance(mySharedMethods(), path='/api')
   # Register dispatchers for single functions
   server.registerFunction(setcookie, path='/api')
   server.registerFunction(echo, path='/api')
   server.registerFunction(block, path='/api')
   server.registerFunction(myip, path='/api')
   server.registerFunction(big, path='/api')
   server.registerFunction(stats, path='/api')
   # Run server
   server.serveForever()
   # Now you can access this api by path http://127.0.0.1:7001/api for JSON-RPC requests
   # Or by path http://127.0.0.1:7001/api/<method>?jsonp=<callback>&(params) for JSONP requests
   #    For example by http://127.0.0.1:7001/api/echo?data=test_data&jsonp=jsonpCallback_129620
<commit_msg>Version 0.4.1 Sample for creating alias<commit_after># -*- coding: utf-8 -*-
import sys, time, random
sys.path.append('/var/python/libs/')
sys.path.append('/var/python/')
sys.path.append('/home/python/libs/')
sys.path.append('/home/python/')

from flaskJSONRPCServer import flaskJSONRPCServer

class mySharedMethods:
   def random(self, mult=65536):
      # Sipmly return random value (0..mult)
      return int(random.random()*mult)

def echo(data='Hello world!'):
   # Simply echo
   return data

def myip(_connection=None):
   # Return client's IP
   return 'Hello, %s!'%(_connection.ip)

def setcookie(_connection=None):
   # Set cookie to client
   _connection.cookiesOut.append({'name':'myTestCookie', 'value':'Your IP is %s'%_connection.ip})
   return 'Setted'

def block():
   # Test for notification request. When it fully implemented, client must  not wait for compliting this function
   time.sleep(10)
   return 'ok'

def stats(_connection=None):
   #calculate connections per second
   tArr1={'connPerSec_now':round(_connection.server.connPerMinute.count/60.0, 2), 'connPerSec_old':round(_connection.server.connPerMinute.oldCount/60.0, 2), 'connPerSec_max':round(_connection.server.connPerMinute.maxCount/60.0, 2), 'speedStats':{}}
   #calculate spped stats
   for k, v in _connection.server.speedStats.items():
      tArr1['speedStats'][k+'_min']=round(min(v), 2)
      tArr1['speedStats'][k+'_max']=round(max(v), 2)
      tArr1['speedStats'][k+'_average']=round(sum(v)/float(len(v)), 2)
   return tArr1

def big(_connection=None):
   _connection.allowCompress=True #allow compression for this method
   s="""
They say that if you sat an infinite number of monkeys with an infinite number of typewriters, then they would write the complete works of Shakespeare. This is known as the Infinite Monkey Theorem. I always feel it would be nice to try the actual experiment out. I mean all these monkeys with typewriters would be kinda cool.

However wouldn't it be easier to get one monkey and one typewriter and give him an infinite amount of time to see how he goes?

Although in theory it sounds simpler, getting hold of even a single monkey is not that easy.

But what if we created a simulated monkey in Python? Now we are talking! That would be easy, and unlike a real monkey our simulated one could work all day and all night. No food required.

Sounds like a plan for a blog post.

However I want to give our virtual monkey a chance. If he managed to type the complete works of Shakespeare without worrying about punctuation, we would be happy right? Same goes for capital letters, let's just ignore them for now. We want to give him a fighting chance of completing his task.

To help you get started I have taken a text file of the complete works of Shakespeare, and removed all punctuation and made all letters small. If any of you are thinking I did that by hand, think again. Python is your friend for automating tasks like that.

Because the complete works is a large tome, I have also included a similar file, with just Hamlet. You can choose which one you want your monkey to attempt!

As with all my blog posts I will give you the complete program first of all. Have a read through this and see how much of it you understand. Try and figure out those sections you don’t. I will then go through the program a line at a time explaining it in more detail.
We do this by using scriptRead.count(monkeyTyped) >= 1 What does that even mean? Well we take the complete works of Shakespeare stored in scriptRead, and we do a count of how many times the thing in the brackets appears. If we put monkeyTyped into the brackets then it will count how many times monkeyCount appeared in scriptRead. We only need it to appear once but if it appears more than once, then that is ok as well. We check this with the greater than or equal to symbol >=.

But the line also has monkeyTyped = ''. Why is that? Well when we start off this while loop the monkey has not typed anything, as we have set the variable to = '' , which means an empty string. This would mean our while loop would be false straight away, and our monkey would never start typing.

To get around this we us an or condition. We check that monkeyTyped is '' or scriptRead.count(monkeyTyped) >= 1

The next line is the line which does the actual typing.
We take monkeyTyped and we add the result of (random.choice(string.ascii_lowercase + ' ')) onto it and save the result as monkeyTyped.

What on earth is (random.choice(string.ascii_lowercase + ' ')) ?

Well we want to simulate the monkey typing on a keyboard. We have already said to make it easy we will not ask for punctuation, although we should ask him for a space between words. Thats not too much to ask is it? We also said we will not be too concerned about it being in capital letters.

So random.choice picks something from inside the brackets by random selection. So what have we put in the brackets. The first thing is

string.ascii_lowercase

This creates a string of the lower case ascii letters. i.e

abcdefghijklmnopqrstuvwxyz

We also add onto the end of that a space using ' '

Now remember this has a space between the two speech-marks. This is different from when we are creating a blank string, which doesn't have a space.

Once he has typed a word, the loop will start again and check if what has been typed is in the complete works. If it is he will continue. If not he will go to the else statement.
Here we return what he has typed. However remember we only get to this place if what is typed is not in the complete works. That's no use to us! We will need to remove the last letter of what he typed, as it was that last letter we know to have made perfect Shakespeare prose into gibberish.

So we return monkeyTyped[:-1] which returns everything up to the last letter but not the last letter.

Ok back to where we were in the main program.

Remember we said we would keep track of the number of key presses? Well lets update that now. We know the monkey has just typed a word, so it is easy to determine the size of that word and increase keyPresses by that amount. Oh but the monkey typed a letter which turned his prose into gibberish. We should add that on as well as we want our count to be accurate!
This line checks to see if the number of key presses he has made is a multiple of 10,000. The % checks if keyPresses / 10,000 has no remainder i.e. it is a multiple of 10,000. If it is then it will print the number of key-presses made. This only works if the key-presses are an exact multiple of 10,000. There will be times when the monkey is in the middle of a word when this happen, so it will not report back. However it reports enough to make you realise the monkey is still working and not fallen asleep.

Remember to press F5 to save and run your program.

That is the end of the program. All you need now it to download the complete works of Shakespeare or just Hamlet if thats what you want to use. Use the links below to do that.
   """
   return s

if __name__=='__main__':
   print 'Running api..'
   # Creating instance of server
   #    <blocking>      set is this server async
   #    <cors>          switch auto CORS support
   #    <gevent>        switch to using Gevent as backend
   #    <debug>         switch to logging connection's info from Flask
   #    <log>           switch to logging debug info from flaskJSONRPCServer
   #    <fallback>      switch auto fallback to JSONP on GET requests
   #    <allowCompress> switch auto compression
   server=flaskJSONRPCServer(("0.0.0.0", 7001), blocking=False, cors=True, gevent=True, debug=False, log=True, fallback=True, allowCompress=False)
   # Register dispatcher for all methods of instance
   server.registerInstance(mySharedMethods(), path='/api')
   # Register dispatchers for single functions
   server.registerFunction(setcookie, path='/api')
   server.registerFunction(echo, path='/api')
   server.registerFunction(block, path='/api')
   server.registerFunction(myip, path='/api')
   big._alias='bigdata' #setting alias for method
   server.registerFunction(big, path='/api')
   server.registerFunction(stats, path='/api')
   # Run server
   server.serveForever()
   # Now you can access this api by path http://127.0.0.1:7001/api for JSON-RPC requests
   # Or by path http://127.0.0.1:7001/api/<method>?jsonp=<callback>&(params) for JSONP requests
   #    For example by http://127.0.0.1:7001/api/echo?data=test_data&jsonp=jsonpCallback_129620
<|endoftext|>


<commit_before>from plugins import CommunicationPlugin
from chat import Message, ChatSession
from socket import socket

class IRCPlugin (CommunicationPlugin):
    def __init__( self, recv_callback, connection_data):
        """
        Initialize a communication plugin.

        :param recv_callback: A function that takes one Message object that should be processed by the main program
        :param args: Further plugin-specific arguments
        :param kwargs: Further plugin-specific arguments
        """
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM);
        self._socket.connect(args[0]);
        self._activeSession = ChatSession('Echo Session');
        self._recv_callback = recv_callback

    def login( credentials : dict ):
        """
        Log in with at the service with credentials

        :param credentials: Data for login
        """
        self._loggedIn = True

    def logout():
        """
        Log out from the service and disconnect
        """
        self._loggedOut = True

    def send(self, message ):
        """
        Send a message

        :param message: Object of type Message with all the required data
        """
        self._recv_callback(message)

    def get_active_sessions() -> [ChatSession]:
        """
        Get a list of all currently active sessions

        :return: A list of all active sessions that self plugin manages
        """
        return self._activeSession

if __name__ == '__main__':
    def test_cb(message):
        print("Received Message: " + message.contact + " " + message.direction + " " + message.message)
    plugin = EchoPlugin(test_cb)
    plugin.send(Message("Testuser", "->", "Hello World!"))
<commit_msg>IRC Plugin<commit_after>from plugins import CommunicationPlugin
from chat import Message, ChatSession
from socket import socket

class IRCPlugin (CommunicationPlugin):
    def __init__( self, recv_callback, connection_data):
        """
        Initialize a communication plugin.

        :param recv_callback: A function that takes one Message object that should be processed by the main program
        :param args: Further plugin-specific arguments
        :param kwargs: Further plugin-specific arguments
        """
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM);
        self._socket.connect(connection_data);
        self._socketfile = self._socket.makefile()
        self._activeSession = ChatSession('Echo Session');
        self._recv_callback = recv_callback

    def login( credentials : dict ):
        """
        Log in with at the service with credentials

        :param credentials: Data for login
        """
        self._socketfile.write('PASS ' + credentials['password'] + '\n')
        self._socketfile.write('NICK ' + credentials['username'] + '\n')
        self._socketfile.write('USER ' + credentials['username'] + ' hostname servername ' + credentials['username'] + '\n')

        self._loggedIn = True

    def logout():
        """
        Log out from the service and disconnect
        """
        self._loggedOut = True

    def send(self, message ):
        """
        Send a message

        :param message: Object of type Message with all the required data
        """
        self._recv_callback(message)

    def get_active_sessions() -> [ChatSession]:
        """
        Get a list of all currently active sessions

        :return: A list of all active sessions that self plugin manages
        """
        return self._activeSession

    def listen():
        pass

if __name__ == '__main__':
    def test_cb(message):
        print("Received Message: " + message.contact + " " + message.direction + " " + message.message)
    plugin = EchoPlugin(test_cb)
    plugin.send(Message("Testuser", "->", "Hello World!"))
<|endoftext|>


<commit_before>from api import app, db
from flask import (render_template, jsonify, request, redirect, url_for,
                   Response, make_response)
from models import Event
from flask.ext.security import login_required
import json
from datetime import datetime

result = [] # Global result list to return result as JSON
e = Event() # Global Instance of Event

@app.route('/api/event/getevents/')
@login_required
def get_events():
    '''
    returns all the events with GET request
    '''
    global result
    allEvents = Event.objects.all()
    create_dict(allEvents)
    return Response(json.dumps(result, cls=PythonJSONEncoder), status=200,
                    content_type="application/json")

@app.errorhandler(404)
def not_found(error):
    return make_response(jsonify({"Error": "Not found"}), 404)

class PythonJSONEncoder(json.JSONEncoder):
    """
    Custom JSON Encoder to encode unsupported data-types to pythonic
    representations.
    """
    def default(self, obj):
        if isinstance(obj, Event):
            return obj.get_dict()
        elif isinstance(obj, datetime):
            return repr(obj.isoformat())
        elif isinstance(obj, datetime.date):
            return repr(obj.isoformat())
        elif isinstance(obj, datetime.time):
            return repr(obj.isoformat())
        else:
        	return repr(obj)
        return super(PythonJSONEncoder, self).default(obj)

def unjsonify(dct):
    if 'eid' in dct:
        dct['eid']=round(eval(dct['eid']), 1)
    return dct

def create_dict(allEvents):
    global result # To store the result of all events
    result = [] # Empty for each call
    for item in allEvents:
        d = {} # To make a dictionary for JSON Response
        d['eid'] = item.eid
        d['title'] = item.title
        d['event_date'] = item.event_date
        d['link'] = item.link
        d['description'] = item.description
        d['venue'] = item.venue
        d['registeration_form_link'] = item.registeration_form_link
        d['event_image_link'] = item.event_image_link
        result.append(d)
    return result<commit_msg>Filteration for past and future events added.<commit_after>from api import app, db
from flask import (render_template, jsonify, request, redirect, url_for,
                   Response, make_response)
from models import Event
from flask.ext.security import login_required
import json
from datetime import datetime

result = [] # Global result list to return result as JSON
e = Event() # Global Instance of Event

@app.route('/api/event/getevents/')
@login_required
def get_events():
    '''
    returns all the events with GET request
    '''
    global result
    time = request.args.get('time')
    if time == "past":
        allEvents = Event.objects.filter(event_date__lte=datetime.now())
    elif time == "future":
        allEvents = Event.objects.filter(event_date__gt=datetime.now())
    else:
        allEvents = Event.objects.all()
    create_dict(allEvents)
    return Response(json.dumps(result, cls=PythonJSONEncoder), status=200,
                    content_type="application/json")

@app.errorhandler(404)
def not_found(error):
    return make_response(jsonify({"Error": "Not found"}), 404)

class PythonJSONEncoder(json.JSONEncoder):
    """
    Custom JSON Encoder to encode unsupported data-types to pythonic
    representations.
    """
    def default(self, obj):
        if isinstance(obj, Event):
            return obj.get_dict()
        elif isinstance(obj, datetime):
            return repr(obj.isoformat())
        elif isinstance(obj, datetime.date):
            return repr(obj.isoformat())
        elif isinstance(obj, datetime.time):
            return repr(obj.isoformat())
        else:
        	return repr(obj)
        return super(PythonJSONEncoder, self).default(obj)

def unjsonify(dct):
    if 'eid' in dct:
        dct['eid']=round(eval(dct['eid']), 1)
    return dct

def create_dict(allEvents):
    global result # To store the result of all events
    result = [] # Empty for each call
    for item in allEvents:
        d = {} # To make a dictionary for JSON Response
        d['eid'] = item.eid
        d['title'] = item.title
        d['event_date'] = item.event_date
        d['link'] = item.link
        d['description'] = item.description
        d['venue'] = item.venue
        d['registeration_form_link'] = item.registeration_form_link
        d['event_image_link'] = item.event_image_link
        result.append(d)
    return result<|endoftext|>


<commit_before>"""simplemarkup, version 2.0, https://github.com/smwst/simplemarkup

Copyright (c) 2012, Ching Chow
All rights reserved.

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

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

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

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

"""

class SimpleMarkup():
    """A simple pretty print markup generator."""
    def __init__(self, width=4):
        self._output = []
        self._open_tags = []
        self._width = " " * width
        self._depth = 0

    def _check(self, value):
        """Raises ValueError if given is not a string."""
        if not isinstance(value, str):
            raise ValueError("must be of type string")
        return value

    def _check_attr(self, attr):
        """Raises ValueError if given is not a list of 2-tuple strings."""
        if not isinstance(attr, list):
            raise ValueError("must be of type list")
        for element in attr:
            if not isinstance(element, tuple) or len(element) != 2:
                raise ValueError("list elements must be a tuple of length 2")
        for name, value in attr:
            if not isinstance(name, str) or not isinstance(value, str):
                raise ValueError("tuple values must be of type string")
        return attr

    def _untrim(self, string, newline):
        """Inserts indentation and newline to a string."""
        lead = "" if not self._output or self._output[-1][-1] != "\n" else self._depth * self._width
        trail = "" if not newline else "\n"
        return lead + string + trail

    def _attr_str(self, attr):
        """Turns a list of attribute tuples into a string."""
        attribute = []
        for name, value in attr:
            attribute.append('{}="{}"'.format(name, value))
        return " " + " ".join(attribute) if len(attribute) > 0 else ""

    def _raw(self, string, newline):
        """Add a string exactly as given."""
        self._output.append(self._untrim(string, newline))
        return self

    def _empty(self, tag, attr, newline):
        """Add an empty element."""
        self._output.append(self._untrim("<{}{} />".format(tag, self._attr_str(attr)), newline))
        return self

    def _begin(self, tag, attr, value, newline):
        """Begin and add a new element."""
        self._open_tags.append(tag)
        self._output.append(self._untrim("<{}{}>{}".format(tag, self._attr_str(attr), value), newline))
        self._depth += 1
        return self

    def _end(self, newline):
        """End an element."""
        self._depth -= 1
        self._output.append(self._untrim("</{}>".format(self._open_tags.pop()), newline))
        return self

    # # # # #

    def output(self):
        """Returns the output in pretty print."""
        return "".join(self._output)

    def raw(self, string):
        """Add a string exactly as given."""
        return self._raw(self._check(string), False)

    def rawln(self, string):
        """Add a string exactly as given and start a newline."""
        return self._raw(self._check(string), True)

    def empty(self, tag, attr=[]):
        """Add an empty element."""
        return self._empty(self._check(tag), self._check_attr(attr), False)

    def emptyln(self, tag, attr=[]):
        """Add an empty element and start a newline."""
        return self._empty(self._check(tag), self._check_attr(attr), True)

    def begin(self, tag, attr=[], value=""):
        """Add the opening tag of an element."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), False)

    def beginln(self, tag, attr=[], value=""):
        """Add the opening tag of an element and start a newline."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), True)

    def end(self):
        """Add the closing tag of an element."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(False)

    def endln(self):
        """Add the closing tag of an element and start a newline."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(True)
<commit_msg>Fixed: Value wasn't placed on a newline when using beginln().<commit_after>"""simplemarkup, version 2.1, https://github.com/smwst/simplemarkup

Copyright (c) 2012, Ching Chow
All rights reserved.

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

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

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

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

"""

class SimpleMarkup():
    """A simple pretty print markup generator."""
    def __init__(self, width=4):
        self._output = []
        self._open_tags = []
        self._width = " " * width
        self._depth = 0

    def _check(self, value):
        """Raises ValueError if given is not a string."""
        if not isinstance(value, str):
            raise ValueError("must be of type string")
        return value

    def _check_attr(self, attr):
        """Raises ValueError if given is not a list of 2-tuple strings."""
        if not isinstance(attr, list):
            raise ValueError("must be of type list")
        for element in attr:
            if not isinstance(element, tuple) or len(element) != 2:
                raise ValueError("list elements must be a tuple of length 2")
        for name, value in attr:
            if not isinstance(name, str) or not isinstance(value, str):
                raise ValueError("tuple values must be of type string")
        return attr

    def _untrim(self, string, newline):
        """Inserts indentation and newline to a string."""
        lead = "" if not self._output or self._output[-1][-1] != "\n" else self._depth * self._width
        trail = "" if not newline else "\n"
        return lead + string + trail

    def _attr_str(self, attr):
        """Turns a list of attribute tuples into a string."""
        attribute = []
        for name, value in attr:
            attribute.append('{}="{}"'.format(name, value))
        return " " + " ".join(attribute) if len(attribute) > 0 else ""

    def _raw(self, string, newline):
        """Add a string exactly as given."""
        self._output.append(self._untrim(string, newline))
        return self

    def _empty(self, tag, attr, newline):
        """Add an empty element."""
        self._output.append(self._untrim("<{}{} />".format(tag, self._attr_str(attr)), newline))
        return self

    def _begin(self, tag, attr, value, newline):
        """Begin and add a new element."""
        self._open_tags.append(tag)
        self._output.append(self._untrim("<{}{}>".format(tag, self._attr_str(attr)), newline))
        self._depth += 1
        if len(value) > 0:
            self._output.append(self._untrim(value, newline))
        return self

    def _end(self, newline):
        """End an element."""
        self._depth -= 1
        self._output.append(self._untrim("</{}>".format(self._open_tags.pop()), newline))
        return self

    # # # # #

    def output(self):
        """Returns the output in pretty print."""
        return "".join(self._output)

    def raw(self, string):
        """Add a string exactly as given."""
        return self._raw(self._check(string), False)

    def rawln(self, string):
        """Add a string exactly as given and start a newline."""
        return self._raw(self._check(string), True)

    def empty(self, tag, attr=[]):
        """Add an empty element."""
        return self._empty(self._check(tag), self._check_attr(attr), False)

    def emptyln(self, tag, attr=[]):
        """Add an empty element and start a newline."""
        return self._empty(self._check(tag), self._check_attr(attr), True)

    def begin(self, tag, attr=[], value=""):
        """Add the opening tag of an element."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), False)

    def beginln(self, tag, attr=[], value=""):
        """Add the opening tag of an element and start a newline."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), True)

    def end(self):
        """Add the closing tag of an element."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(False)

    def endln(self):
        """Add the closing tag of an element and start a newline."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(True)
<|endoftext|>


<commit_before>import logging

from neo4j_util import Neo4JException
import neo4j_util as db_util

log = logging.getLogger('rhizi')

class DB_Driver_Base():

    def log_committed_queries(self, statement_set):
        for sp_dict in statement_set['statements']:
            if None != sp_dict['parameters']:
                msg = '\tq: {0}\n\tp: {1}'.format(sp_dict['statement'].encode('utf-8'),
                                                  sp_dict['parameters'])
            else:
                msg = '\tq: {0}'.format(sp_dict['statement'])
            log.debug(msg)

class DB_Driver_Embedded(DB_Driver_Base):
    def __init__(self, db_base_url):
        self.tx_base_url = db_base_url + '/db/data/transaction'

        from org.rhizi.db.neo4j.util import EmbeddedNeo4j
        self.edb = EmbeddedNeo4j.createDb()
        self.edb.createDb()

    def begin_tx(self, op):
        pass

    def exec_query_set(self, op):
        pass

    def commit_tx(self, op):
        pass

class DB_Driver_REST(DB_Driver_Base):
    def __init__(self, db_base_url):
        self.tx_base_url = db_base_url + '/db/data/transaction'

    def begin_tx(self, op):
        tx_open_url = self.tx_base_url

        try:
            #
            # [!] neo4j seems picky about receiving an additional empty statement list
            #
            data = db_util.db_query_set_to_REST_form([])  # open TX with empty query_set
            ret = db_util.post_neo4j(tx_open_url, data)
            tx_commit_url = ret['commit']
            op.parse_tx_id(tx_commit_url)

            log.debug('tx-open: id: {0}, commit-url: {1}'.format(op.tx_id, tx_commit_url))
        except Exception as e:
            raise Exception('failed to open transaction: ' + e.message)

    def exec_query_set(self, op):

        tx_url = "{0}/{1}".format(self.tx_base_url, op.tx_id)
        q_set = db_util.db_query_set_to_REST_form(op.query_set)

        try:
            post_ret = db_util.post_neo4j(tx_url, q_set)
            op.result_set = post_ret['results']
            op.error_set = post_ret['errors']
            if 0 != len(op.error_set):
                raise Neo4JException(op.error_set)

            self.log_committed_queries(q_set)
        except Neo4JException as e:
            # NOTE: python 2.7 loses the stack when reraising the exception.
            # python 3 does the right thing, but gevent doesn't support it yet.
            log.error('REST statement: %r' % statement_set)
            log.exception(e)
            raise e
        except Exception as e:
            log.error('REST statement: %r' % statement_set)
            log.exception(e)
            raise Exception('failed exec op statements: err: {0}, url: {1}'.format(e.message, tx_url))

    def commit_tx(self, op):
        tx_commit_url = "{0}/{1}/commit".format(self.tx_base_url, op.tx_id)

        try:
            #
            # [!] neo4j seems picky about receiving an additional empty statement list
            #
            data = db_util.db_query_set_to_REST_form([])  # close TX with empty query_set
            ret = db_util.post(tx_commit_url, data)

            log.debug('tx-commit: id: {0}, commit-url: {1}'.format(op.tx_id, tx_commit_url))

            return ret
        except Exception as e:
            raise Exception('failed to commit transaction:' + e.message)
<commit_msg>misc rename: q_set -> statement_param_pair_set<commit_after>import logging

from neo4j_util import Neo4JException
import neo4j_util as db_util

log = logging.getLogger('rhizi')

class DB_Driver_Base():

    def log_committed_queries(self, statement_set):
        for sp_dict in statement_set['statements']:
            if None != sp_dict['parameters']:
                msg = '\tq: {0}\n\tp: {1}'.format(sp_dict['statement'].encode('utf-8'),
                                                  sp_dict['parameters'])
            else:
                msg = '\tq: {0}'.format(sp_dict['statement'])
            log.debug(msg)

class DB_Driver_Embedded(DB_Driver_Base):

    def __init__(self, db_base_url):
        self.tx_base_url = db_base_url + '/db/data/transaction'

        from org.rhizi.db.neo4j.util import EmbeddedNeo4j
        self.edb = EmbeddedNeo4j.createDb()
        self.edb.createDb()

    def begin_tx(self, op):
        pass

    def exec_query_set(self, op):
        pass

    def commit_tx(self, op):
        pass

class DB_Driver_REST(DB_Driver_Base):

    def __init__(self, db_base_url):
        self.tx_base_url = db_base_url + '/db/data/transaction'

    def begin_tx(self, op):
        tx_open_url = self.tx_base_url

        try:
            #
            # [!] neo4j seems picky about receiving an additional empty statement list
            #
            data = db_util.db_query_set_to_REST_form([])  # open TX with empty query_set
            ret = db_util.post_neo4j(tx_open_url, data)
            tx_commit_url = ret['commit']
            op.parse_tx_id(tx_commit_url)

            log.debug('tx-open: id: {0}, commit-url: {1}'.format(op.tx_id, tx_commit_url))
        except Exception as e:
            raise Exception('failed to open transaction: ' + e.message)

    def exec_query_set(self, op):

        tx_url = "{0}/{1}".format(self.tx_base_url, op.tx_id)
        statement_param_pair_set = db_util.db_query_set_to_REST_form(op.query_set)

        try:
            post_ret = db_util.post_neo4j(tx_url, statement_param_pair_set)
            op.result_set = post_ret['results']
            op.error_set = post_ret['errors']
            if 0 != len(op.error_set):
                raise Neo4JException(op.error_set)

            self.log_committed_queries(statement_param_pair_set)
        except Neo4JException as e:
            # NOTE: python 2.7 loses the stack when reraising the exception.
            # python 3 does the right thing, but gevent doesn't support it yet.
            log.error('REST statement: %r' % statement_param_pair_set)
            log.exception(e)
            raise e
        except Exception as e:
            log.error('REST statement: %r' % statement_param_pair_set)
            log.exception(e)
            raise Exception('failed exec op statements: err: {0}, url: {1}'.format(e.message, tx_url))

    def commit_tx(self, op):
        tx_commit_url = "{0}/{1}/commit".format(self.tx_base_url, op.tx_id)

        try:
            #
            # [!] neo4j seems picky about receiving an additional empty statement list
            #
            data = db_util.db_query_set_to_REST_form([])  # close TX with empty query_set
            ret = db_util.post(tx_commit_url, data)

            log.debug('tx-commit: id: {0}, commit-url: {1}'.format(op.tx_id, tx_commit_url))

            return ret
        except Exception as e:
            raise Exception('failed to commit transaction:' + e.message)
<|endoftext|>


<commit_before>class SimpleMarkup():
    """A simple pretty print markup generator."""
    def __init__(self, width=4):
        self._output = []
        self._open_tags = []
        self._width = " " * width
        self._depth = 0

    def _check(self, value):
        """Raises ValueError if given is not a string."""
        if not isinstance(value, str):
            raise ValueError("must be of type string")
        return value

    def _check_attr(self, attr):
        """Raises ValueError if given is not a list of 2-tuple strings."""
        if not isinstance(attr, list):
            raise ValueError("must be of type list")
        for element in attr:
            if not isinstance(element, tuple) or len(element) != 2:
                raise ValueError("list elements must be a tuple of length 2")
        for name, value in attr:
            if not isinstance(name, str) or not isinstance(value, str):
                raise ValueError("tuple values must be of type string")
        return attr

    def _untrim(self, string, newline):
        """Inserts indentation and newline to a string."""
        lead = "" if not self._output or self._output[-1][-1] != "\n" else self._depth * self._width
        trail = "" if not newline else "\n"
        return lead + string + trail

    def _attr_str(self, attr):
        """Turns a list of attribute tuples into a string."""
        attribute = []
        for name, value in attr:
            attribute.append('{}="{}"'.format(name, value))
        return " " + " ".join(attribute) if len(attribute) > 0 else ""

    def _raw(self, string, newline):
        """Add a string exactly as given."""
        self._output.append(self._untrim(string, newline))
        return self

    def _empty(self, tag, attr, newline):
        """Add an empty element."""
        self._output.append(self._untrim("<{}{} />".format(tag, self._attr_str(attr)), newline))
        return self

    def _begin(self, tag, attr, value, newline):
        """Begin and add a new element."""
        self._open_tags.append(tag)
        self._output.append(self._untrim("<{}{}>{}".format(tag, self._attr_str(attr), value), newline))
        self._depth += 1
        return self

    def _end(self, newline):
        """End an element."""
        self._depth -= 1
        self._output.append(self._untrim("</{}>".format(self._open_tags.pop()), newline))
        return self

    # # # # #

    def output(self):
        """Returns the output in pretty print."""
        return "".join(self._output)

    def raw(self, string):
        """Add a string exactly as given."""
        return self._raw(self._check(string), False)

    def rawln(self, string):
        """Add a string exactly as given and start a newline."""
        return self._raw(self._check(string), True)

    def empty(self, tag, attr=[]):
        """Add an empty element."""
        return self._empty(self._check(tag), self._check_attr(attr), False)

    def emptyln(self, tag, attr=[]):
        """Add an empty element and start a newline."""
        return self._empty(self._check(tag), self._check_attr(attr), True)

    def begin(self, tag, attr=[], value=""):
        """Begin and add a new element."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), False)

    def beginln(self, tag, attr=[], value=""):
        """Begin, add a new element, and start a newline."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), True)

    def end(self):
        """End an element."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(False)

    def endln(self):
        """End an element and start a newline."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(True)
<commit_msg>Use updated simplemarkup.<commit_after>"""simplemarkup, version 2.1, https://github.com/smwst/simplemarkup

Copyright (c) 2012, Ching Chow
All rights reserved.

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

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

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

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

"""

class SimpleMarkup():
    """A simple pretty print markup generator."""
    def __init__(self, width=4):
        self._output = []
        self._open_tags = []
        self._width = " " * width
        self._depth = 0

    def _check(self, value):
        """Raises ValueError if given is not a string."""
        if not isinstance(value, str):
            raise ValueError("must be of type string")
        return value

    def _check_attr(self, attr):
        """Raises ValueError if given is not a list of 2-tuple strings."""
        if not isinstance(attr, list):
            raise ValueError("must be of type list")
        for element in attr:
            if not isinstance(element, tuple) or len(element) != 2:
                raise ValueError("list elements must be a tuple of length 2")
        for name, value in attr:
            if not isinstance(name, str) or not isinstance(value, str):
                raise ValueError("tuple values must be of type string")
        return attr

    def _untrim(self, string, newline):
        """Inserts indentation and newline to a string."""
        lead = "" if not self._output or self._output[-1][-1] != "\n" else self._depth * self._width
        trail = "" if not newline else "\n"
        return lead + string + trail

    def _attr_str(self, attr):
        """Turns a list of attribute tuples into a string."""
        attribute = []
        for name, value in attr:
            attribute.append('{}="{}"'.format(name, value))
        return " " + " ".join(attribute) if len(attribute) > 0 else ""

    def _raw(self, string, newline):
        """Add a string exactly as given."""
        self._output.append(self._untrim(string, newline))
        return self

    def _empty(self, tag, attr, newline):
        """Add an empty element."""
        self._output.append(self._untrim("<{}{} />".format(tag, self._attr_str(attr)), newline))
        return self

    def _begin(self, tag, attr, value, newline):
        """Begin and add a new element."""
        self._open_tags.append(tag)
        self._output.append(self._untrim("<{}{}>".format(tag, self._attr_str(attr)), newline))
        self._depth += 1
        if len(value) > 0:
            self._output.append(self._untrim(value, newline))
        return self

    def _end(self, newline):
        """End an element."""
        self._depth -= 1
        self._output.append(self._untrim("</{}>".format(self._open_tags.pop()), newline))
        return self

    # # # # #

    def output(self):
        """Returns the output in pretty print."""
        return "".join(self._output)

    def raw(self, string):
        """Add a string exactly as given."""
        return self._raw(self._check(string), False)

    def rawln(self, string):
        """Add a string exactly as given and start a newline."""
        return self._raw(self._check(string), True)

    def empty(self, tag, attr=[]):
        """Add an empty element."""
        return self._empty(self._check(tag), self._check_attr(attr), False)

    def emptyln(self, tag, attr=[]):
        """Add an empty element and start a newline."""
        return self._empty(self._check(tag), self._check_attr(attr), True)

    def begin(self, tag, attr=[], value=""):
        """Add the opening tag of an element."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), False)

    def beginln(self, tag, attr=[], value=""):
        """Add the opening tag of an element and start a newline."""
        return self._begin(self._check(tag), self._check_attr(attr), self._check(value), True)

    def end(self):
        """Add the closing tag of an element."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(False)

    def endln(self):
        """Add the closing tag of an element and start a newline."""
        if self._depth == 0:
            raise IndexError("already ended")
        return self._end(True)
<|endoftext|>


<commit_before>import logging

import claripy

from .s_action_object import SimActionObject

l = logging.getLogger("simuvex.s_cc")

class SimFunctionArgument(object):
    def __init__(self):
        pass

class SimRegArg(SimFunctionArgument):
    def __init__(self, reg_name):
        SimFunctionArgument.__init__(self)

        self.name = reg_name

    def __repr__(self):
        return "<%s>" % self.name

class SimStackArg(SimFunctionArgument):
    def __init__(self, stack_offset):
        SimFunctionArgument.__init__(self)

        self.offset = stack_offset

    def __repr__(self):
        return "[%xh]" % self.offset

class SimCC(object):
    '''
    This is the base class for all calling conventions. You should not directly instantiate this class.
    '''
    RET_VAL_REG = None
    ARG_REGS = None
    STACKARG_SP_DIFF = None

    def __init__(self, arch, sp_delta=None):
        self.arch = arch
        self.sp_delta = sp_delta

        # A list of argument positions
        self.args = None
        # A list of return value positions
        self.ret_vals = None

    def set_args(self, state, args):
        """
        Sets the value @expr as being the @index-th argument of a function
        """

        # Normalize types of all arguments
        bv_args = [ ]
        for expr in args:
            if type(expr) in (int, long):
                e = state.BVV(expr, state.arch.bits)
            elif type(expr) in (str,):
                e = state.BVV(expr)
            elif not isinstance(expr, (claripy.A, SimActionObject)):
                raise SimCCError("can't set argument of type %s" % type(expr))
            else:
                e = expr

            if len(e) != state.arch.bits:
                raise SimCCError("all args must be %d bits long" % state.arch.bits)

            bv_args.append(e)

        reg_offsets = self.ARG_REGS
        if reg_offsets is None:
            raise NotImplementedError('ARG_REGS is not specified for calling convention %s' % type(self))

        if len(args) > len(reg_offsets):
            stack_shift = (len(args) - len(reg_offsets)) * state.arch.stack_change
            sp_value = state.reg_expr('sp') + stack_shift
            state.store_reg('sp', sp_value)
        else:
            sp_value = state.reg_expr('sp')

        for index, e in reversed(tuple(enumerate(bv_args))):
            self.arg_setter(state, e, reg_offsets, sp_value, index)

    # Returns a bitvector expression representing the nth argument of a function
    def arg(self, state, index, stackarg_mem_base=None):
        reg_offsets = self.ARG_REGS
        if reg_offsets is None:
            raise NotImplementedError('ARG_REGS is not specified for calling convention %s' % type(self))

        if self.STACKARG_SP_DIFF is None:
            raise NotImplementedError('STACKARG_SP_DIFF is not specified for calling convention %s' % type(self))

        if stackarg_mem_base is None:
            # This is the default case, which is used inside SimProcedures.
            stackarg_mem_base = state.reg_expr(self.arch.sp_offset) + self.STACKARG_SP_DIFF

        return self.arg_getter(state, reg_offsets, stackarg_mem_base, index)

    # Sets an expression as the return value. Also updates state.
    def set_return_expr(self, state, expr):
        expr = self._normalize_return_expr(state, expr)

        if self.RET_VAL_REG is None:
            raise NotImplementedError('RET_VAL_REG is not specified for calling convention %s' % type(self))

        state.store_reg(self.RET_VAL_REG, expr)

    def _normalize_return_expr(self, state, expr):
        if type(expr) in (int, long):
            expr = state.BVV(expr, state.arch.bits)

        return expr

    #
    # Helper functions
    #

    # Helper function to get an argument, given a list of register locations it can be and stack information for overflows.
    def arg_getter(self, state, reg_offsets, args_mem_base, index):
        stack_step = -state.arch.stack_change

        if index < len(reg_offsets):
            expr = state.reg_expr(reg_offsets[index], endness=state.arch.register_endness)
        else:
            index -= len(reg_offsets)
            mem_addr = args_mem_base + (index * stack_step)
            expr = state.mem_expr(mem_addr, stack_step, endness=state.arch.memory_endness)

        return expr

    def arg_setter(self, state, expr, reg_offsets, args_mem_base, index):
        stack_step = -state.arch.stack_change

        # Set register parameters
        if index < len(reg_offsets):
            offs = reg_offsets[index]
            state.store_reg(offs, expr, endness=state.arch.register_endness)

        # Set remaining parameters on the stack
        else:
            index -= len(reg_offsets)
            mem_addr = args_mem_base + (index * stack_step)
            state.store_mem(mem_addr, expr, endness=state.arch.memory_endness)

    @staticmethod
    def match(project, function_address, cfg=None):
        '''
        Try to decide the arguments to this function.
        `cfg` is not necessary, but providing a CFG makes our life easier and will give you a better analysis
        result (i.e. we have an idea of how this function is called in its call-sites).
        If a CFG is not provided or we cannot find the given function address in the given CFG, we will generate
        a local CFG of the function to detect how it is using the arguments.
        '''
        arch = project.arch

        args = [ ]
        ret_vals = [ ]
        sp_delta = 0

        #
        # Determine how many arguments this function has.
        #
        func = cfg.function_manager.function(function_address)
        if func is not None:
            reg_args, stack_args = func.arguments

            for arg in reg_args:
                a = SimRegArg(project.arch.register_names[arg])
                args.append(a)

            for arg in stack_args:
                a = SimStackArg(arg)
                args.append(a)

            sp_delta = func.sp_delta

            for c in CC:
                if c._match(project, args, sp_delta):
                    return c(project.arch, args, ret_vals, sp_delta)

        else:
            # TODO:
            pass

        # We cannot determine the calling convention of this function.

        return SimCCUnknown(arch, args, ret_vals, sp_delta)

    @property
    def arguments(self):
        return self.args

    def __repr__(self):
        return "SimCC"

class SimCCCdecl(SimCC):
    ARG_REGS = [ ] # All arguments are passed in stack
    STACKARG_SP_DIFF = 4 # Return address is pushed on to stack by call
    RET_VAL_REG = 'eax'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimX86 and sp_delta == 0:
            any_reg_args = any([a for a in args if isinstance(a, SimStackArg)])

            if not any_reg_args:
                return True

        return False

class SimCCSystemVAMD64(SimCC):
    ARG_REGS = [ 'rdi', 'rsi', 'rdx', 'rcx', 'r8', 'r9' ]
    STACKARG_SP_DIFF = 8 # Return address is pushed on to stack by call
    RET_VAL_REG = 'rax'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

        # Remove the ret address on stack
        if self.args is not None:
            self.args = [ i for i in self.args if not (isinstance(i, SimStackArg) and i.offset == 0x8) ]

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimAMD64 and sp_delta == 0:
            reg_args = [ i.name for i in args if isinstance(i, SimRegArg)]
            for r in SimCCSystemVAMD64.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # There is still something left.
                # Bad!
                return False

            stack_args = [ i.offset for i in args if isinstance(i, SimStackArg) ]
            if 0x8 not in stack_args:
                # Where is the return address?
                # TODO: Are we too strict about this?
                return False

            # That's it!
            return True

        return False

    def __repr__(self):
        return "System V AMD64 - %s %s" % (self.arch.name, self.args)

class SimCCARM(SimCC):
    ARG_REGS = [ 'r0', 'r1', 'r2', 'r3' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'r0'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimARM and sp_delta == 0:
            reg_args = [ i.name for i in args if isinstance(i, SimRegArg) ]

            for r in SimCCARM.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCO32(SimCC):
    ARG_REGS = [ 'a0', 'a1', 'a2', 'a3' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'v0'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimMIPS32 and sp_delta == 0:
            reg_args = [i.name for i in args if isinstance(i, SimRegArg)]

            for r in SimCCO32.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCPowerPC(SimCC):
    ARG_REGS = [ 'r3', 'r4', 'r5', 'r6', 'r7', 'r8', 'r9', 'r10' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'r3'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimPPC32 and sp_delta == 0:
            reg_args = [i.name for i in args if isinstance(i, SimRegArg)]

            for r in SimCCPowerPC.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCPowerPC64(SimCC):
    ARG_REGS = [ 'r3', 'r4', 'r5', 'r6', 'r7', 'r8', 'r9', 'r10' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'r3'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimPPC64 and sp_delta == 0:
            reg_args = [i.name for i in args if isinstance(i, SimRegArg)]

            for r in SimCCPowerPC64.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCUnknown(SimCC):
    '''
    WOW an unknown calling convention!
    '''
    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args
        self.ret_vals = ret_vals

    def arg_reg_offsets(self):
        pass

    @staticmethod
    def _match(p, args, sp_delta):

        # It always returns True
        return True

    def __repr__(self):
        s = "UnknownCC - %s %s sp_delta=%d" % (self.arch.name, self.args, self.sp_delta)
        return s

CC = [ SimCCCdecl, SimCCSystemVAMD64, SimCCARM, SimCCO32, SimCCPowerPC, SimCCPowerPC64 ]
DefaultCC = {
    'AMD64': SimCCSystemVAMD64,
    'X86': SimCCCdecl,
    'ARM': SimCCARM,
    'MIPS32': SimCCO32,
    'PPC32': SimCCPowerPC,
    'PPC64': SimCCPowerPC64,
}

from .s_errors import SimCCError
from .s_arch import SimX86, SimAMD64, SimARM, SimMIPS32, SimPPC32, SimPPC64
<commit_msg>make SimCC work with typedast<commit_after>import logging

import claripy

from .s_action_object import SimActionObject

l = logging.getLogger("simuvex.s_cc")

class SimFunctionArgument(object):
    def __init__(self):
        pass

class SimRegArg(SimFunctionArgument):
    def __init__(self, reg_name):
        SimFunctionArgument.__init__(self)

        self.name = reg_name

    def __repr__(self):
        return "<%s>" % self.name

class SimStackArg(SimFunctionArgument):
    def __init__(self, stack_offset):
        SimFunctionArgument.__init__(self)

        self.offset = stack_offset

    def __repr__(self):
        return "[%xh]" % self.offset

class SimCC(object):
    '''
    This is the base class for all calling conventions. You should not directly instantiate this class.
    '''
    RET_VAL_REG = None
    ARG_REGS = None
    STACKARG_SP_DIFF = None

    def __init__(self, arch, sp_delta=None):
        self.arch = arch
        self.sp_delta = sp_delta

        # A list of argument positions
        self.args = None
        # A list of return value positions
        self.ret_vals = None

    def set_args(self, state, args):
        """
        Sets the value @expr as being the @index-th argument of a function
        """

        # Normalize types of all arguments
        bv_args = [ ]
        for expr in args:
            if type(expr) in (int, long):
                e = state.BVV(expr, state.arch.bits)
            elif type(expr) in (str,):
                e = state.BVV(expr)
            elif not isinstance(expr, (claripy.Base, SimActionObject)):
                raise SimCCError("can't set argument of type %s" % type(expr))
            else:
                e = expr

            if len(e) != state.arch.bits:
                raise SimCCError("all args must be %d bits long" % state.arch.bits)

            bv_args.append(e)

        reg_offsets = self.ARG_REGS
        if reg_offsets is None:
            raise NotImplementedError('ARG_REGS is not specified for calling convention %s' % type(self))

        if len(args) > len(reg_offsets):
            stack_shift = (len(args) - len(reg_offsets)) * state.arch.stack_change
            sp_value = state.reg_expr('sp') + stack_shift
            state.store_reg('sp', sp_value)
        else:
            sp_value = state.reg_expr('sp')

        for index, e in reversed(tuple(enumerate(bv_args))):
            self.arg_setter(state, e, reg_offsets, sp_value, index)

    # Returns a bitvector expression representing the nth argument of a function
    def arg(self, state, index, stackarg_mem_base=None):
        reg_offsets = self.ARG_REGS
        if reg_offsets is None:
            raise NotImplementedError('ARG_REGS is not specified for calling convention %s' % type(self))

        if self.STACKARG_SP_DIFF is None:
            raise NotImplementedError('STACKARG_SP_DIFF is not specified for calling convention %s' % type(self))

        if stackarg_mem_base is None:
            # This is the default case, which is used inside SimProcedures.
            stackarg_mem_base = state.reg_expr(self.arch.sp_offset) + self.STACKARG_SP_DIFF

        return self.arg_getter(state, reg_offsets, stackarg_mem_base, index)

    # Sets an expression as the return value. Also updates state.
    def set_return_expr(self, state, expr):
        expr = self._normalize_return_expr(state, expr)

        if self.RET_VAL_REG is None:
            raise NotImplementedError('RET_VAL_REG is not specified for calling convention %s' % type(self))

        state.store_reg(self.RET_VAL_REG, expr)

    def _normalize_return_expr(self, state, expr):
        if type(expr) in (int, long):
            expr = state.BVV(expr, state.arch.bits)

        return expr

    #
    # Helper functions
    #

    # Helper function to get an argument, given a list of register locations it can be and stack information for overflows.
    def arg_getter(self, state, reg_offsets, args_mem_base, index):
        stack_step = -state.arch.stack_change

        if index < len(reg_offsets):
            expr = state.reg_expr(reg_offsets[index], endness=state.arch.register_endness)
        else:
            index -= len(reg_offsets)
            mem_addr = args_mem_base + (index * stack_step)
            expr = state.mem_expr(mem_addr, stack_step, endness=state.arch.memory_endness)

        return expr

    def arg_setter(self, state, expr, reg_offsets, args_mem_base, index):
        stack_step = -state.arch.stack_change

        # Set register parameters
        if index < len(reg_offsets):
            offs = reg_offsets[index]
            state.store_reg(offs, expr, endness=state.arch.register_endness)

        # Set remaining parameters on the stack
        else:
            index -= len(reg_offsets)
            mem_addr = args_mem_base + (index * stack_step)
            state.store_mem(mem_addr, expr, endness=state.arch.memory_endness)

    @staticmethod
    def match(project, function_address, cfg=None):
        '''
        Try to decide the arguments to this function.
        `cfg` is not necessary, but providing a CFG makes our life easier and will give you a better analysis
        result (i.e. we have an idea of how this function is called in its call-sites).
        If a CFG is not provided or we cannot find the given function address in the given CFG, we will generate
        a local CFG of the function to detect how it is using the arguments.
        '''
        arch = project.arch

        args = [ ]
        ret_vals = [ ]
        sp_delta = 0

        #
        # Determine how many arguments this function has.
        #
        func = cfg.function_manager.function(function_address)
        if func is not None:
            reg_args, stack_args = func.arguments

            for arg in reg_args:
                a = SimRegArg(project.arch.register_names[arg])
                args.append(a)

            for arg in stack_args:
                a = SimStackArg(arg)
                args.append(a)

            sp_delta = func.sp_delta

            for c in CC:
                if c._match(project, args, sp_delta):
                    return c(project.arch, args, ret_vals, sp_delta)

        else:
            # TODO:
            pass

        # We cannot determine the calling convention of this function.

        return SimCCUnknown(arch, args, ret_vals, sp_delta)

    @property
    def arguments(self):
        return self.args

    def __repr__(self):
        return "SimCC"

class SimCCCdecl(SimCC):
    ARG_REGS = [ ] # All arguments are passed in stack
    STACKARG_SP_DIFF = 4 # Return address is pushed on to stack by call
    RET_VAL_REG = 'eax'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimX86 and sp_delta == 0:
            any_reg_args = any([a for a in args if isinstance(a, SimStackArg)])

            if not any_reg_args:
                return True

        return False

class SimCCSystemVAMD64(SimCC):
    ARG_REGS = [ 'rdi', 'rsi', 'rdx', 'rcx', 'r8', 'r9' ]
    STACKARG_SP_DIFF = 8 # Return address is pushed on to stack by call
    RET_VAL_REG = 'rax'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

        # Remove the ret address on stack
        if self.args is not None:
            self.args = [ i for i in self.args if not (isinstance(i, SimStackArg) and i.offset == 0x8) ]

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimAMD64 and sp_delta == 0:
            reg_args = [ i.name for i in args if isinstance(i, SimRegArg)]
            for r in SimCCSystemVAMD64.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # There is still something left.
                # Bad!
                return False

            stack_args = [ i.offset for i in args if isinstance(i, SimStackArg) ]
            if 0x8 not in stack_args:
                # Where is the return address?
                # TODO: Are we too strict about this?
                return False

            # That's it!
            return True

        return False

    def __repr__(self):
        return "System V AMD64 - %s %s" % (self.arch.name, self.args)

class SimCCARM(SimCC):
    ARG_REGS = [ 'r0', 'r1', 'r2', 'r3' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'r0'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimARM and sp_delta == 0:
            reg_args = [ i.name for i in args if isinstance(i, SimRegArg) ]

            for r in SimCCARM.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCO32(SimCC):
    ARG_REGS = [ 'a0', 'a1', 'a2', 'a3' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'v0'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimMIPS32 and sp_delta == 0:
            reg_args = [i.name for i in args if isinstance(i, SimRegArg)]

            for r in SimCCO32.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCPowerPC(SimCC):
    ARG_REGS = [ 'r3', 'r4', 'r5', 'r6', 'r7', 'r8', 'r9', 'r10' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'r3'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimPPC32 and sp_delta == 0:
            reg_args = [i.name for i in args if isinstance(i, SimRegArg)]

            for r in SimCCPowerPC.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCPowerPC64(SimCC):
    ARG_REGS = [ 'r3', 'r4', 'r5', 'r6', 'r7', 'r8', 'r9', 'r10' ]
    STACKARG_SP_DIFF = 0
    RET_VAL_REG = 'r3'

    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args

    @staticmethod
    def _match(p, args, sp_delta):
        if type(p.arch) is SimPPC64 and sp_delta == 0:
            reg_args = [i.name for i in args if isinstance(i, SimRegArg)]

            for r in SimCCPowerPC64.ARG_REGS:
                if r in reg_args:
                    reg_args.remove(r)
            if reg_args:
                # Still something left...
                return False

            return True

        return False

class SimCCUnknown(SimCC):
    '''
    WOW an unknown calling convention!
    '''
    def __init__(self, arch, args=None, ret_vals=None, sp_delta=None):
        SimCC.__init__(self, arch, sp_delta)

        self.args = args
        self.ret_vals = ret_vals

    def arg_reg_offsets(self):
        pass

    @staticmethod
    def _match(p, args, sp_delta):

        # It always returns True
        return True

    def __repr__(self):
        s = "UnknownCC - %s %s sp_delta=%d" % (self.arch.name, self.args, self.sp_delta)
        return s

CC = [ SimCCCdecl, SimCCSystemVAMD64, SimCCARM, SimCCO32, SimCCPowerPC, SimCCPowerPC64 ]
DefaultCC = {
    'AMD64': SimCCSystemVAMD64,
    'X86': SimCCCdecl,
    'ARM': SimCCARM,
    'MIPS32': SimCCO32,
    'PPC32': SimCCPowerPC,
    'PPC64': SimCCPowerPC64,
}

from .s_errors import SimCCError
from .s_arch import SimX86, SimAMD64, SimARM, SimMIPS32, SimPPC32, SimPPC64
<|endoftext|>


<commit_before>from __future__ import print_function

import numpy as np

# Import sile objects
from sisl.io.sile import *

# Import the geometry object
from sisl import Geometry, Atom, SuperCell, Grid
from sisl.unit import unit_convert

__all__ = ['CUBESile']

Ang2Bohr = unit_convert('Ang', 'Bohr')

class CUBESile(Sile):
    """ CUBE file object """

    def _open(self):
        if self.file.endswith('gz'):
            self.fh = gzip.open(self.file)
        else:
            # Default to buffer 1MB
            self.fh = open(self.file, self._mode, buffering=1024**2)
        self._line = 0

    @Sile_fh_open
    def write_geometry(self, geom, size=None,
            fmt='15.10e', origo=None,
            *args, **kwargs):
        """ Writes the `Geometry` object attached to this grid """
        sile_raise_write(self)

        # Write header
        self._write('\n')
        self._write('sisl --- CUBE file\n')

        if size is None:
            size = np.ones([3], np.int32)
        if origo is None:
            origo = np.zeros([3], np.float64)

        _fmt = '{:d} {:15.10e} {:15.10e} {:15.10e}\n'

        # Add #-of atoms and origo
        self._write(_fmt.format(len(geom), *origo * Ang2Bohr))

        # Write the cell and voxels
        dcell = np.empty([3, 3], np.float64)
        for ix in range(3):
            dcell[ix, :] = geom.cell[ix, :] / size[ix] * Ang2Bohr
        self._write(_fmt.format(size[0], *dcell[0, :]))
        self._write(_fmt.format(size[1], *dcell[1, :]))
        self._write(_fmt.format(size[2], *dcell[2, :]))

        tmp = ' {:' + fmt + '}'
        _fmt = '{:d} 0.0' + tmp + tmp + tmp + '\n'
        for ia in geom:
            self._write(_fmt.format(geom.atom[ia].Z, *geom.xyz[ia, :] * Ang2Bohr))

    @Sile_fh_open
    def write_grid(self, grid, fmt='%.5e', *args, **kwargs):
        """ Writes the geometry to the contained file """
        # Check that we can write to the file
        sile_raise_write(self)

        # Write the geometry
        self.write_geometry(grid.geom, size=grid.grid.shape, *args, **kwargs)

        # A CUBE file contains grid-points aligned like this:
        # for x
        #   for y
        #     for z
        #       write...
        cast = np.vectorize((fmt + '\n').__mod__)
        for z in np.nditer(np.asarray(grid.grid, order='C'), flags=['external_loop', 'buffered'],
                           op_flags=[['readonly']], order='C', buffersize=grid.shape[2]):
            self._write(cast(z))

        # Add a finishing line to ensure empty ending
        self._write('\n')

    @Sile_fh_open
    def read_supercell(self, na=False):
        """ Returns `SuperCell` object from the CUBE file

        If `na=True` it will return a tuple (na,SuperCell)
        """

        self.readline()  # header 1
        self.readline()  # header 2
        tmp = self.readline().split()  # origo
        na = int(tmp[0])

        cell = np.empty([3, 3], np.float64)
        for i in [0, 1, 2]:
            tmp = self.readline().split()
            s = int(tmp[0])
            tmp = tmp[1:]
            for j in [0, 1, 2]:
                cell[i, j] = float(tmp[j]) * s

        cell = cell / Ang2Bohr
        if na:
            return na, SuperCell(cell)
        return SuperCell(cell)

    @Sile_fh_open
    def read_geometry(self):
        """ Returns `Geometry` object from the CUBE file """
        na, sc = self.read_supercell(na=True)

        # Start reading the geometry
        xyz = np.empty([na, 3], np.float64)
        atom = []
        for ia in range(na):
            tmp = self.readline().split()
            atom.append(Atom(int(tmp[0])))
            xyz[ia, 0] = float(tmp[2])
            xyz[ia, 1] = float(tmp[3])
            xyz[ia, 2] = float(tmp[4])

        xyz /= Ang2Bohr
        return Geometry(xyz, atom, sc=sc)

    @Sile_fh_open
    def read_grid(self):
        """ Returns `Grid` object from the CUBE file """
        geom = self.read_geometry()

        # Now seek behind to read grid sizes
        self.fh.seek(0)

        # Skip headers and origo
        self.readline()
        self.readline()
        na = int(self.readline().split()[0])

        ngrid = [0] * 3
        for i in [0, 1, 2]:
            tmp = self.readline().split()
            ngrid[i] = int(tmp[0])

        # Read past the atoms
        for i in range(na):
            self.readline()

        grid = Grid(ngrid, dtype=np.float32, geom=geom)
        grid.grid = np.loadtxt(self.fh, dtype=grid.dtype)
        grid.grid.shape = ngrid

        return grid

add_sile('cube', CUBESile, case=False, gzip=True)
<commit_msg>bug: final fix for removal of savetxt<commit_after>from __future__ import print_function

import numpy as np

# Import sile objects
from sisl.io.sile import *

# Import the geometry object
from sisl import Geometry, Atom, SuperCell, Grid
from sisl.unit import unit_convert

__all__ = ['CUBESile']

Ang2Bohr = unit_convert('Ang', 'Bohr')

class CUBESile(Sile):
    """ CUBE file object """

    def _open(self):
        if self.file.endswith('gz'):
            self.fh = gzip.open(self.file)
        else:
            # Default to buffer 1MB
            self.fh = open(self.file, self._mode, buffering=1024**2)
        self._line = 0

    @Sile_fh_open
    def write_geometry(self, geom, size=None,
            fmt='15.10e', origo=None,
            *args, **kwargs):
        """ Writes the `Geometry` object attached to this grid """
        sile_raise_write(self)

        # Write header
        self._write('\n')
        self._write('sisl --- CUBE file\n')

        if size is None:
            size = np.ones([3], np.int32)
        if origo is None:
            origo = np.zeros([3], np.float64)

        _fmt = '{:d} {:15.10e} {:15.10e} {:15.10e}\n'

        # Add #-of atoms and origo
        self._write(_fmt.format(len(geom), *origo * Ang2Bohr))

        # Write the cell and voxels
        dcell = np.empty([3, 3], np.float64)
        for ix in range(3):
            dcell[ix, :] = geom.cell[ix, :] / size[ix] * Ang2Bohr
        self._write(_fmt.format(size[0], *dcell[0, :]))
        self._write(_fmt.format(size[1], *dcell[1, :]))
        self._write(_fmt.format(size[2], *dcell[2, :]))

        tmp = ' {:' + fmt + '}'
        _fmt = '{:d} 0.0' + tmp + tmp + tmp + '\n'
        for ia in geom:
            self._write(_fmt.format(geom.atom[ia].Z, *geom.xyz[ia, :] * Ang2Bohr))

    @Sile_fh_open
    def write_grid(self, grid, fmt='%.5e', *args, **kwargs):
        """ Writes the geometry to the contained file """
        # Check that we can write to the file
        sile_raise_write(self)

        # Write the geometry
        self.write_geometry(grid.geom, size=grid.grid.shape, *args, **kwargs)

        # A CUBE file contains grid-points aligned like this:
        # for x
        #   for y
        #     for z
        #       write...
        cast = np.vectorize((fmt + '\n').__mod__)
        _fmt = (fmt + '\n') * grid.shape[2]
        for z in np.nditer(np.asarray(grid.grid, order='C'), flags=['external_loop', 'buffered'],
                           op_flags=[['readonly']], order='C', buffersize=grid.shape[2]):
            self._write(_fmt % tuple(z))

        # Add a finishing line to ensure empty ending
        self._write('\n')

    @Sile_fh_open
    def read_supercell(self, na=False):
        """ Returns `SuperCell` object from the CUBE file

        If `na=True` it will return a tuple (na,SuperCell)
        """

        self.readline()  # header 1
        self.readline()  # header 2
        tmp = self.readline().split()  # origo
        na = int(tmp[0])

        cell = np.empty([3, 3], np.float64)
        for i in [0, 1, 2]:
            tmp = self.readline().split()
            s = int(tmp[0])
            tmp = tmp[1:]
            for j in [0, 1, 2]:
                cell[i, j] = float(tmp[j]) * s

        cell = cell / Ang2Bohr
        if na:
            return na, SuperCell(cell)
        return SuperCell(cell)

    @Sile_fh_open
    def read_geometry(self):
        """ Returns `Geometry` object from the CUBE file """
        na, sc = self.read_supercell(na=True)

        # Start reading the geometry
        xyz = np.empty([na, 3], np.float64)
        atom = []
        for ia in range(na):
            tmp = self.readline().split()
            atom.append(Atom(int(tmp[0])))
            xyz[ia, 0] = float(tmp[2])
            xyz[ia, 1] = float(tmp[3])
            xyz[ia, 2] = float(tmp[4])

        xyz /= Ang2Bohr
        return Geometry(xyz, atom, sc=sc)

    @Sile_fh_open
    def read_grid(self):
        """ Returns `Grid` object from the CUBE file """
        geom = self.read_geometry()

        # Now seek behind to read grid sizes
        self.fh.seek(0)

        # Skip headers and origo
        self.readline()
        self.readline()
        na = int(self.readline().split()[0])

        ngrid = [0] * 3
        for i in [0, 1, 2]:
            tmp = self.readline().split()
            ngrid[i] = int(tmp[0])

        # Read past the atoms
        for i in range(na):
            self.readline()

        grid = Grid(ngrid, dtype=np.float32, geom=geom)
        grid.grid = np.loadtxt(self.fh, dtype=grid.dtype)
        grid.grid.shape = ngrid

        return grid

add_sile('cube', CUBESile, case=False, gzip=True)
<|endoftext|>