Initial commit

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;
+}

bitkey_proto/build.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+cd `dirname $0`
+
+protoc --python_out=. bitkey.proto

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()

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)

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)

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

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