def test_register_frame(self):
self.assertRaises(ValueError, frame.RegisterFrame, 'ab')
self.assertFalse(frame.registration._frametypes.get('x', False))
@frame.RegisterFrame('x')
class MyXFrame(frame.base.BaseFrame):
pass
self.assertEqual(frame.registration._frametypes[RawData('x')],
MyXFrame)
self.assertRaises(ValueError, frame.RegisterFrame, 'x')
decorator = frame.RegisterFrame('y')
class MyYFrame(frame.base.BaseFrame):
pass
decorator(MyYFrame)
self.assertEqual(frame.registration._frametypes[RawData('y')],
MyYFrame)
decorator(MyXFrame)
self.assertEqual(frame.registration._frametypes[RawData('y')],
MyYFrame)
def test_create_class():
@frame.RegisterFrame('q')
class NotAFrame(object):
pass
self.assertRaises(TypeError, test_create_class)
def construct(identity, content):
return EncryptedFrame(
RawData('r') + \
RawData(str_address(identity.location)) + \
RawData((0,)) + \
RawData(identity.encryptor.encrypt(content))
)
def test_decompress(self):
compressed = bz2.compress(RawData('foobar').export())
self.assertEqual(
frame.compressed.Bz2Compressor(compressed).decompress(),
RawData('foobar'))
self.assertRaises(frame.compressed.CompressionError,
frame.compressed.Bz2Compressor('foobar').decompress)
def test_index(self):
r = RawData("abc")
self.assertEqual(r.index(RawData("b")), 1)
self.assertEqual(r.index(99), 2)
self.assertRaises(TypeError, r.index, RawData("ab"))
self.assertRaises(TypeError, r.index, 300)
self.assertRaises(ValueError, r.index, RawData("d"))
def construct(identity, content):
signature = identity.sign(content)
siglen = len(signature)
return SignedFrame(
RawData('s') + \
RawData(str_address(identity.location)) + \
RawData((0, siglen // 256, siglen % 256)) + \
RawData(signature) + \
RawData(content)
)
def construct(content, compression_type):
c_type = String(compression_type)
if c_type in _compression_alias:
c_type = _compression_alias[c_type]
elif c_type.toRawData() in _compression_types:
c_type = c_type.toRawData()
else:
raise ValueError('unknown compression_type')
return CompressedFrame(
RawData('c') + c_type + RawData(0) + \
RawData(_compression_types[c_type](content).compress())
)
def test_hash(self):
self.assertEqual(
self.re.hash('hello, world'),
RawData((0x09, 0xca, 0x7e, 0x4e, 0xaa, 0x6e, 0x8a, 0xe9, 0xc7,
0xd2, 0x61, 0x16, 0x71, 0x29, 0x18, 0x48, 0x83, 0x64,
0x4d, 0x07, 0xdf, 0xba, 0x7c, 0xbf, 0xbc, 0x4c, 0x8a,
0x2e, 0x08, 0x36, 0x0d, 0x5b)))
self.assertEqual(
self.re.hash('a' * 300),
RawData((0x98, 0x35, 0xfa, 0x6b, 0xf4, 0xe2, 0x0a, 0x9b, 0x9e,
0xa8, 0x12, 0x50, 0x63, 0x02, 0xe9, 0x89, 0x82, 0x72,
0x1a, 0x6c, 0xf8, 0xd2, 0xca, 0xe6, 0x7a, 0xf5, 0x71,
0x29, 0xbf, 0x21, 0xae, 0x90)))
def test_sign(self):
sig = self.re.sign(self.plaintext)
self.assertEqual(
sig,
RawData((0xe3, 0xa4, 0x58, 0x28, 0x84, 0x48, 0x64, 0xc3, 0xa1,
0xac, 0x3b, 0xf0, 0x4b, 0x03, 0xf2, 0x22, 0x5d, 0x3e,
0x27, 0xe1, 0xb9, 0x94, 0x56, 0x99, 0x96, 0x26, 0x64,
0x08, 0xe2, 0x10, 0xe7, 0x35)))
self.assertEqual(
self.re.encrypt(sig),
RawData((0x09, 0xca, 0x7e, 0x4e, 0xaa, 0x6e, 0x8a, 0xe9, 0xc7,
0xd2, 0x61, 0x16, 0x71, 0x29, 0x18, 0x48, 0x83, 0x64,
0x4d, 0x07, 0xdf, 0xba, 0x7c, 0xbf, 0xbc, 0x4c, 0x8a,
0x2e, 0x08, 0x36, 0x0d, 0x5b)))
self.assertEqual(self.re.encrypt(sig), self.re.hash(self.plaintext))
def test_generate(self):
myrsa = RSA(None, 1024)
proto = myrsa.proto()
self.assertEqual(len(proto), 2)
self.assertEqual(proto[0], 'rsa')
self.assertTrue(proto[1].startswith(
RawData('-----BEGIN RSA PRIVATE KEY-----\n').export()))
def test_decode(self):
from ejtp.identity import Identity, IdentityCache
cache = IdentityCache()
cache[['testing']] = Identity('joe', ['rotate', 1], ['testing'])
self.assertEqual(
frame.encrypted.EncryptedFrame('r["testing"]\x00gpp').decode(
cache), RawData('foo'))
def test_decode(self):
self.assertEqual(
frame.compressed.construct('foobar', 'z').decode(),
RawData('foobar'))
self.assertRaises(
ValueError,
frame.compressed.CompressedFrame('cunknown\x00foobar').decode)
def test_decode(self):
from ejtp.identity import Identity, IdentityCache
ident = Identity('joe', ['rotate', 1], ['testing'])
cache = IdentityCache()
cache[['testing']] = ident
signature = ident.sign('foo')
siglen = len(signature)
signed_content = RawData('s["testing"]\x00') + (
siglen // 256, siglen % 256) + signature + 'foo'
self.assertEqual(
frame.signed.SignedFrame(signed_content).decode(cache),
RawData('foo'))
self.assertRaises(
ValueError,
frame.signed.SignedFrame(
's["testing"]\x00\x00\x07invalidfoo').decode, cache)
def message(self, target, mhash):
self.write_json(
target,
{
'type': 'ejforward-message',
'target': target,
'data': RawData(
self.client(target)['messages'][String(mhash)])
},
)
def test_construct(self):
from ejtp.identity import Identity
ident = Identity('joe', ['rotate', 1], ['testing'])
signature = ident.sign('foo')
siglen = len(signature)
self.assertEqual(
frame.signed.construct(ident, 'foo'),
frame.signed.SignedFrame(
RawData('s["testing"]\x00') + (siglen // 256, siglen % 256) +
signature + 'foo'))
def test_init(self):
self.assertEqual(
frame.base.BaseFrame('foobar')._content, RawData('foobar'))
self.assertRaises(TypeError, frame.base.BaseFrame, 1234)
ancestor = frame.base.BaseFrame('oldfoobar')
self.assertEqual(
frame.base.BaseFrame('foobar', [ancestor])._ancestors,
[ancestor.crop()])
self.assertRaises(TypeError, frame.base.BaseFrame, 'foobar', 123)
self.assertRaises(TypeError, frame.base.BaseFrame, 'foobar',
[object()])
def test_construct(self):
for comp_t, comp in frame.compressed._compression_types.items():
compressed = comp('foobar').compress()
self.assertEqual(
frame.compressed.construct('foobar', comp_t),
frame.compressed.CompressedFrame(
RawData('c') + comp_t + '\x00' + compressed))
for comp_alias, comp_t in frame.compressed._compression_alias.items():
self.assertEqual(frame.compressed.construct('foobar', comp_alias),
frame.compressed.construct('foobar', comp_t))
self.assertRaises(ValueError, frame.compressed.construct, 'foobar',
'unkknown')
def inject(self, newdata):
'''
Process new data from the outside world.
'''
self._buffer += newdata
if RawData(".") not in self._buffer:
return
size, content = self._buffer.split('.', 1)
size = int(size.export(), 16) # Read size as hex
if len(content) < size:
return
if self.jack:
self.jack.recv(content[:size])
else:
self._outqueue.put(content[:size])
self._buffer = content[size:]
def rcv_callback(self, msg, client_obj):
data = msg.unpack()
mtype = data['type']
if mtype == 'ejforward-notify':
self._status = data
for callback in self._status_callbacks:
callback(self)
self._status_callbacks = []
elif mtype == 'ejforward-message':
internal = RawData(data['data'])
self.ack([hashfunc(internal)])
try:
self.send(frame.createFrame(internal)) # forward to router
except ValueError:
logger.warning("Invalid frame, discarding")
else:
logger.warning("Unknown message type, %r" % mtype)
def __init__(self, jack=None):
self.jack = jack
self._buffer = RawData()
self._running = False
self._outqueue = Queue.Queue()
self._thread = threading.Thread(target=self.run)
class Connection(object):
'''
Represents a persistent connection to a remote host. Should be subclassed.
'''
def __init__(self, jack=None):
self.jack = jack
self._buffer = RawData()
self._running = False
self._outqueue = Queue.Queue()
self._thread = threading.Thread(target=self.run)
# SUBCLASS INTERFACE ------------------------------------------------------
@RawDataDecorator(strict=True)
def _send(self, frame):
'''
Subclass-provided function that sends text to the remote end.
'''
pass
@RawDataDecorator(args=False, ret=True, strict=True)
def _recv(self):
'''
Subclass-provided function that returns text. Allowed to block.
'''
pass
def run(self):
'''
Network recieve loop.
'''
pass
# PROVIDED BY BASE CLASS --------------------------------------------------
def start(self):
self._running = True
self._thread.start()
def close(self):
self._running = False
def send(self, frame):
self._send(self.wrap(frame.content))
def recv(self, timeout=0):
'''
Retrieve from output queue without blocking. Returns None or str.
The output queue is only used if self.jack == None.
'''
try:
return self._outqueue.get(timeout=timeout)
except Queue.Empty:
return None
def wrap(self, frame):
'''
Wrap a frame for transport
'''
return RawData(format(len(frame), '0x')) + "." + frame
@RawDataDecorator(strict=True)
def inject(self, newdata):
'''
Process new data from the outside world.
'''
self._buffer += newdata
if RawData(".") not in self._buffer:
return
size, content = self._buffer.split('.', 1)
size = int(size.export(), 16) # Read size as hex
if len(content) < size:
return
if self.jack:
self.jack.recv(content[:size])
else:
self._outqueue.put(content[:size])
self._buffer = content[size:]
def wrap(self, frame):
'''
Wrap a frame for transport
'''
return RawData(format(len(frame), '0x')) + "." + frame
class Bz2Compressor(Compressor):
@RawDataDecorator(args=False, ret=True, strict=True)
def compress(self):
return bz2.compress(self._data.export())
@RawDataDecorator(args=False, ret=True, strict=True)
def decompress(self):
try:
return bz2.decompress(self._data.export())
except IOError:
raise CompressionError()
_compression_types = {
RawData('z'): ZlibCompressor,
RawData('b'): Bz2Compressor
}
@RegisterFrame('c')
class CompressedFrame(BaseFrame):
@RawDataDecorator(args=False, ret=True, strict=True)
def decode(self, ident_cache=None):
if self.header not in _compression_types:
raise ValueError('unknown compression type')
compressor = _compression_types[self.header]
return compressor(self.body).decompress()
_compression_alias = {
def test_compress(self):
self.assertEqual(
frame.compressed.Bz2Compressor('foobar').compress(),
RawData(bz2.compress(RawData('foobar').export())))
def test_RawData(self):
self._assert('"test"', RawData('test'))
def encrypt(self, value):
# Uses custom format to encrypt arbitrary length strings with padding
code = RawData(len(value)) + "\x00" + value
code += (16 - len(code) % 16) * "\x00"
return self.cipher.encrypt(code.export())
def _assert(self, expected, value):
value = hasher.strict(value)
self.assertEqual(RawData(expected), RawData(value))
def rotate(self, source, offset):
# not checking args here, because offset would be converted to RawData
result = RawData()
for i in source:
result += (int(i)+offset) % 256
return result
def decrypt(self, value):
code = RawData(self.cipher.decrypt(value.export()))
split = code.index('\x00')
length = int(code[:split])
code = code[split + 1:]
return code[:length]
def construct(content):
return JSONFrame(
RawData('j\x00') + \
RawData(strict(content))
)
def test_init(self):
self.assertEqual(
frame.compressed.Compressor('foobar')._data, RawData('foobar'))
def _assert(self, expected, value):
value = hasher.strictify(json.dumps(value))
self.assertEqual(RawData(expected), RawData(value))