research/2008-displacement/main-cpushare

#! /usr/bin/python

import os, struct, sys, random

import gd

from cpushare.proto_const import *
from cpushare.exceptions import CompilationError

# Our input data
source = False
result = False
todo = 0

# Parse command line
if len(sys.argv) < 4:
    print 'Usage: %s <buy_order.cpu> -<mode> [image list]' % (sys.argv[0],)
    sys.exit(-1)
mode = sys.argv[2][1]
source = sys.argv[3:]
result = [False] * len(source)

class buy_state_machine_class(object):
    buy_api = '0.0.0'

    w, h = 0, 0
    index = -1
    answer = ''

    def __init__(self, protocol):
        self.protocol = protocol
        self.handler = self.stringReceived

    def start(self):
        global todo, source
        if todo >= len(source):
            return # We're finished... FIXME: a process is stuck
        self.index = todo
        todo += 1
        gd.gdMaxColors = 256 * 256 * 256
        im = gd.image(source[self.index])
        # Send mode
        self.protocol.sendString(PROTO_SECCOMP_FORWARD + mode)
        # Send image size
        (self.w, self.h) = (w, h) = im.size()
        print 'sending %s (%dx%d)' % (source[self.index], w, h)
        msg = chr(w / 256) + chr(w % 256) + chr(h / 256) + chr(h % 256) 
        self.protocol.sendString(PROTO_SECCOMP_FORWARD + msg)
        # Send image contents
        msg = ''
        for y in range(h):
            for x in range(w):
                p = im.getPixel((x, y))
                rgb = im.colorComponents(p)
                msg += chr(rgb[1]) # use green coordinate
        self.protocol.sendString(PROTO_SECCOMP_FORWARD + msg)
        print '%s bytes sent' % len(msg)

    def stringReceived(self, string):
        control = string[0]
        data = string[1:]

        if control == PROTO_SECCOMP_FORWARD:
            self.answer += data
            if ord(self.answer[-1]) != 0:
                return
            global result
            result[self.index] = self.answer[:-1]
            print self.answer[:-1],
            self.protocol.sendString(PROTO_SECCOMP_SUCCESS)
            # the PROTO_SECCOMP_SUCCESS probably won't have a chance
            # to go out but it doesn't matter because we'll disconnect
            # not just this but all other seccomp state machines too
            from twisted.internet import reactor
            reactor.stop()
        elif control == PROTO_SECCOMP_SIGNAL:
            print 'Checkpoint starting'
        elif control == PROTO_LOG:
            print repr(data)
        else:
            if control == PROTO_SECCOMP_SUCCESS:
                pass
            elif control == PROTO_SECCOMP_FAILURE:
                status = struct.unpack('!i', data)[0]
                exit_code = status >> 8
                signal = status & 0xff
                s = 'Seccomp failure: status %d, exit_code %d, signal %d.' % \
                      (status, exit_code, signal)
                print s
            else:
                s = 'Unknown failure %d - %s' % (ord(control), repr(string))
                self.protocol.sendString(PROTO_SECCOMP_FAILURE + s)
                print s
            self.protocol.sendString(PROTO_SECCOMP_FAILURE)
            self.protocol.transport.loseConnection()

    def connectionLost(self):
        pass

# Build local module
from new import module
m = module('cpushare_buy')

# Create bytecode -- i686 only
from cpushare.seccomp_gen import seccomp_gen_class
m.seccomp_gen_hash = {}
m.seccomp_gen_hash['i686'] = seccomp_gen_class('bytecode', 'i686')
m.seccomp_gen_hash['i686'].heap_kbytes = 1
m.seccomp_gen_hash['i686'].stack_kbytes = 1

# Estimate of the max number of seconds that the sell client
# will take to checkpoint and send its state back to us
m.checkpoint_sec = 10

# Our buying state machine
m.buy_state_machine_class = buy_state_machine_class

# Append our module to the global list of modules
sys.modules['cpushare_buy'] = m

# Build a new command line and run twistd
sys.argv = [sys.argv[0], '-q', '-n', 'cpushare', '--order', sys.argv[1]]
from twisted.scripts.twistd import run
run()