From 99d892e7592a8f20ff0b73fd078bf42e5a780bbc Mon Sep 17 00:00:00 2001 From: slush Date: Tue, 13 Nov 2012 14:09:39 +0000 Subject: [PATCH] Initial commit --- __init__.py | 0 bitkey_proto/__init__.py | 0 bitkey_proto/bitkey.proto | 147 ++++++++++++++++++++++++++++++++++++++ bitkey_proto/build.sh | 5 ++ bitkey_proto/mapping.py | 51 +++++++++++++ test.py | 70 ++++++++++++++++++ transport.py | 62 ++++++++++++++++ transport_pipe.py | 54 ++++++++++++++ transport_serial.py | 35 +++++++++ 9 files changed, 424 insertions(+) create mode 100644 __init__.py create mode 100644 bitkey_proto/__init__.py create mode 100644 bitkey_proto/bitkey.proto create mode 100755 bitkey_proto/build.sh create mode 100644 bitkey_proto/mapping.py create mode 100755 test.py create mode 100644 transport.py create mode 100644 transport_pipe.py create mode 100644 transport_serial.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bitkey_proto/__init__.py b/bitkey_proto/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bitkey_proto/bitkey.proto b/bitkey_proto/bitkey.proto new file mode 100644 index 0000000..212661d --- /dev/null +++ b/bitkey_proto/bitkey.proto @@ -0,0 +1,147 @@ +enum Algorithm { + BIP32 = 0; + ELECTRUM = 1; +} + +enum ScriptType { + PAYTOADDRESS = 0; + PAYTOSCRIPTHASH = 1; +} + +// Response: None or Features +message Initialize { +} + +message Features { + optional string version = 1; + optional bool otp = 2; + optional bool pin = 3; + optional bool spv = 4; + repeated Algorithm algo = 5; +} + +// Description: Test if another side is still alive. +// Response: None or Success +message Ping { +} + +// Description: Response message for previous request with given id. +message Success { + optional string message = 1; +} + +// Description: Response message for previous request with given id. +message Failure { + optional int32 code = 1; + optional string message = 2; +} + +// Response: UUID or Failure +message GetUUID { +} + +message UUID { + required bytes UUID = 1; +} + +message OtpRequest { + optional string message = 1; +} + +message OtpAck { + required string otp = 1; +} + +message OtpCancel { +} + +message PinRequest { + optional string message = 1; +} + +message PinAck { + required string pin = 1; +} + +message PinCancel { +} + +// Response: OtpRequest, Entropy, Failure +message GetEntropy { + required uint32 size = 1; +} + +message Entropy { + required bytes entropy = 1; +} + +// Response: MasterPublicKey, Failure +message GetMasterPublicKey { + required Algorithm algo = 1 [default=BIP32]; +} + +message MasterPublicKey { + required bytes key = 1; +} + +// Response: Success, OtpRequest, Failure +message LoadDevice { + required string seed = 1; + optional bool otp = 2 [default=true]; + optional string pin = 3; + optional bool spv = 4 [default=true]; +} + +// Response: Success, OtpRequest, PinRequest, Failure +message ResetDevice { +} + +message TxOutput { + required string address = 1; + repeated uint32 address_n = 2; + required uint64 amount = 3; + required ScriptType script_type = 4; + repeated bytes script_args = 5; +} + +// Response: Success, SignedInput, Failure +message TxInput { + repeated uint32 address_n = 1; + required uint64 amount = 2; + required bytes prev_hash = 3; + required uint32 prev_index = 4; + optional bytes script_sig = 5; +} + +// Response: SignedTx, Success, OtpRequest, PinRequest, Failure +message SignTx { + required Algorithm algo = 1 [default=BIP32]; + optional bool stream = 2; // enable streaming + required uint64 fee = 3; + repeated TxOutput outputs = 4; + repeated TxInput inputs = 5; + optional uint32 inputs_count = 6; // for streaming + optional bytes random = 7; +} + +message SignedTx { + repeated bytes signature = 1; +} + +/* +inputs = [] # list of TxInput +for i in inputs: + for x in inputs: + send(x) + + signature = send(SignInput(i)) +*/ + +// Response: SignedInput, Failure +message SignInput { + required TxInput input = 1; +} + +message SignedInput { + required bytes signature = 1; +} \ No newline at end of file diff --git a/bitkey_proto/build.sh b/bitkey_proto/build.sh new file mode 100755 index 0000000..b96dd45 --- /dev/null +++ b/bitkey_proto/build.sh @@ -0,0 +1,5 @@ +#!/bin/bash + +cd `dirname $0` + +protoc --python_out=. bitkey.proto diff --git a/bitkey_proto/mapping.py b/bitkey_proto/mapping.py new file mode 100644 index 0000000..bf6eee8 --- /dev/null +++ b/bitkey_proto/mapping.py @@ -0,0 +1,51 @@ +import bitkey_pb2 as proto + +map_type_to_class = { + 0: proto.Initialize, + 1: proto.Ping, + 2: proto.Success, + 3: proto.Failure, + 4: proto.GetUUID, + 5: proto.UUID, + 6: proto.OtpRequest, + 7: proto.OtpAck, + 8: proto.OtpCancel, + 9: proto.GetEntropy, + 10: proto.Entropy, + 11: proto.GetMasterPublicKey, + 12: proto.MasterPublicKey, + 13: proto.LoadDevice, + 14: proto.ResetDevice, + 15: proto.SignTx, + 16: proto.SignedTx, + 17: proto.Features, + 18: proto.PinRequest, + 19: proto.PinAck, + 20: proto.PinCancel, +} + +map_class_to_type = {} + +def get_type(msg): + return map_class_to_type[msg.__class__] + +def get_class(t): + return map_type_to_class[t] + +def build_index(): + for k, v in map_type_to_class.items(): + map_class_to_type[v] = k + +def check_missing(): + from google.protobuf import reflection + + types = [ proto.__dict__[item] for item in dir(proto) + if issubclass(proto.__dict__[item].__class__, reflection.GeneratedProtocolMessageType) ] + + missing = list(set(types) - set(map_type_to_class.values())) + + if len(missing): + raise Exception("Following protobuf messages are not defined in mapping: %s" % missing) + +check_missing() +build_index() diff --git a/test.py b/test.py new file mode 100755 index 0000000..3b5be2b --- /dev/null +++ b/test.py @@ -0,0 +1,70 @@ +#!/usr/bin/python + +import time + +from transport_pipe import PipeTransport +from transport_serial import SerialTransport +from bitkey_proto import bitkey_pb2 as proto + +def pprint(msg): + return "<%s>:\n%s" % (msg.__class__.__name__, msg) + +def call(msg, tries=3): + print '----------------------' + print "Sending", pprint(msg) + d.write(msg) + resp = d.read() + + if isinstance(resp, proto.OtpRequest): + if resp.message: + print "Message:", resp.message + otp = raw_input("OTP required: ") + d.write(proto.OtpAck(otp=otp)) + resp = d.read() + + if isinstance(resp, proto.PinRequest): + if resp.message: + print "Message:", resp.message + pin = raw_input("PIN required: ") + d.write(proto.PinAck(pin=pin)) + resp = d.read() + + if isinstance(resp, proto.Failure) and resp.code in (3, 6): + if tries <= 1 and resp.code == 3: + raise Exception("OTP is invalid, too many retries") + if tries <= 1 and resp.code == 6: + raise Exception("PIN is invalid, too many retries") + + # Invalid OTP or PIN, try again + if resp.code == 3: + print "OTP is invalid, let's try again..." + elif resp.code == 6: + print "PIN is invalid, let's try again..." + + return call(msg, tries-1) + + if isinstance(resp, proto.Failure): + raise Exception(resp.code, resp.message) + + print "Received", pprint(resp) + return resp + +d = PipeTransport('../../bitkey-python/device.socket', is_device=False) +#d = SerialTransport('../../bitkey-python/COM9') + +#start = time.time() + +#for x in range(1000): + +call(proto.Initialize()) +call(proto.Ping()) +call(proto.GetUUID()) +#call(proto.GetEntropy(size=10)) +#call(proto.LoadDevice(seed='beyond neighbor scratch swirl embarrass doll cause also stick softly physical nice', +# otp=True, pin='1234', spv=True)) + +#call(proto.ResetDevice()) +call(proto.GetMasterPublicKey(algo=proto.ELECTRUM)) +#call(proto.ResetDevice()) + +#print 10000 / (time.time() - start) diff --git a/transport.py b/transport.py new file mode 100644 index 0000000..5bc83f8 --- /dev/null +++ b/transport.py @@ -0,0 +1,62 @@ +import struct +from bitkey_proto import bitkey_pb2 as proto +from bitkey_proto import mapping + +class Transport(object): + def __init__(self, device, *args, **kwargs): + self.device = device + self._open() + + def _open(self): + raise NotImplemented + + def _close(self): + raise NotImplemented + + def _write(self, msg): + raise NotImplemented + + def _read(self): + raise NotImplemented + + def close(self): + self._close() + + def write(self, msg): + ser = msg.SerializeToString() + header = struct.pack(">HL", mapping.get_type(msg), len(ser)) + self._write("##%s%s" % (header, ser)) + + def read(self): + (msg_type, data) = self._read() + inst = mapping.get_class(msg_type)() + inst.ParseFromString(data) + return inst + + def _read_headers(self, read_f): + # Try to read headers until some sane value are detected + is_ok = False + while not is_ok: + + # Align cursor to the beginning of the header ("##") + c = read_f.read(1) + while c != '#': + if c == '': + # timeout + raise Exception("Timed out while waiting for the magic character") + print "Warning: Aligning to magic characters" + c = read_f.read(1) + + if read_f.read(1) != "#": + # Second character must be # to be valid header + raise Exception("Second magic character is broken") + + # Now we're most likely on the beginning of the header + try: + headerlen = struct.calcsize(">HL") + (msg_type, datalen) = struct.unpack(">HL", read_f.read(headerlen)) + break + except: + raise Exception("Cannot parse header length") + + return (msg_type, datalen) \ No newline at end of file diff --git a/transport_pipe.py b/transport_pipe.py new file mode 100644 index 0000000..2ca1394 --- /dev/null +++ b/transport_pipe.py @@ -0,0 +1,54 @@ +'''TransportFake implements fake wire transport over local named pipe. +Use this transport for talking with bitkey simulator.''' + +import os + +from transport import Transport + +class PipeTransport(Transport): + def __init__(self, device, is_device, *args, **kwargs): + self.is_device = is_device # Set True if act as device + + super(PipeTransport, self).__init__(device, *args, **kwargs) + + def _open(self): + if self.is_device: + self.filename_read = self.device+'.to' + self.filename_write = self.device+'.from' + + os.mkfifo(self.filename_read, 0600) + os.mkfifo(self.filename_write, 0600) + else: + self.filename_read = self.device+'.from' + self.filename_write = self.device+'.to' + + if not os.path.exists(self.filename_write): + raise Exception("Not connected") + + self.write_fd = os.open(self.filename_write, os.O_RDWR)#|os.O_NONBLOCK) + self.write_f = os.fdopen(self.write_fd, 'w+') + + self.read_fd = os.open(self.filename_read, os.O_RDWR)#|os.O_NONBLOCK) + self.read_f = os.fdopen(self.read_fd, 'rb') + + def _close(self): + self.read_f.close() + self.write_f.close() + os.unlink(self.filename_read) + os.unlink(self.filename_write) + + def _write(self, msg): + try: + self.write_f.write(msg) + self.write_f.flush() + except OSError: + print "Error while writing to socket" + raise + + def _read(self): + try: + (msg_type, datalen) = self._read_headers(self.read_f) + return (msg_type, self.read_f.read(datalen)) + except IOError: + print "Failed to read from device" + raise \ No newline at end of file diff --git a/transport_serial.py b/transport_serial.py new file mode 100644 index 0000000..8820096 --- /dev/null +++ b/transport_serial.py @@ -0,0 +1,35 @@ +'''SerialTransport implements wire transport over serial port.''' + +# Local serial port loopback: socat PTY,link=COM8 PTY,link=COM9 + +import serial + +from transport import Transport + +class SerialTransport(Transport): + def __init__(self, device, *args, **kwargs): + self.serial = None + super(SerialTransport, self).__init__(device, *args, **kwargs) + + def _open(self): + self.serial = serial.Serial(self.device, 115200, timeout=10, writeTimeout=10) + + def _close(self): + self.serial.close() + self.serial = None + + def _write(self, msg): + try: + self.serial.write(msg) + self.serial.flush() + except serial.SerialException: + print "Error while writing to socket" + raise + + def _read(self): + try: + (msg_type, datalen) = self._read_headers(self.serial) + return (msg_type, self.serial.read(datalen)) + except serial.SerialException: + print "Failed to read from device" + raise \ No newline at end of file