make some things more general; allow command line arguments

This commit is contained in:
Seth Schoen
2012-07-14 17:31:43 -07:00
parent 1fd5ae1c9d
commit 19bcb8486c
+24 -14
View File
@@ -15,12 +15,22 @@ difficulty = 20
def sha256(m): def sha256(m):
return hashlib.sha256(m).hexdigest() return hashlib.sha256(m).hexdigest()
try: assert len(sys.argv) > 1 or "CHOCOLATESERVER" in os.environ, "Must specify server via command line or CHOCOLATESERVER environment variable."
upstream = "https://%s/chocolate.py" % os.environ["CHOCOLATESERVER"] if len(sys.argv) > 1:
except KeyError: server = sys.argv[1]
print "Please set the environment variable CHOCOLATESERVER to the hostname" else:
print "of a server that speaks this protocol." server = os.environ["CHOCOLATESERVER"]
sys.exit(1)
upstream = "https://%s/chocolate.py" % server
if len(sys.argv) > 3:
req_file = sys.argv[2]
key_file = sys.argv[3]
else:
req_file = "req.pem"
key_file = "key.pem"
cert_file = "cert.pem" # we should use getopt to set all of these
def do(m): def do(m):
u = urllib2.urlopen(upstream, m.SerializeToString()) u = urllib2.urlopen(upstream, m.SerializeToString())
@@ -34,21 +44,20 @@ def init(m):
m.session = "" m.session = ""
def make_request(m, csr): def make_request(m, csr):
server = os.environ["CHOCOLATESERVER"]
m.request.recipient = server m.request.recipient = server
m.request.timestamp = int(time.time()) m.request.timestamp = int(time.time())
m.request.csr = csr m.request.csr = csr
m.request.clientpuzzle = hashcash.mint(server, difficulty) m.request.clientpuzzle = hashcash.mint(server, difficulty)
def sign(k, m): def sign(key, m):
m.request.sig = CSR.sign(k, ("(%d) (%s) (%s)" % (m.request.timestamp, m.request.recipient, m.request.csr))) m.request.sig = CSR.sign(key, ("(%d) (%s) (%s)" % (m.request.timestamp, m.request.recipient, m.request.csr)))
k=chocolatemessage() k=chocolatemessage()
m=chocolatemessage() m=chocolatemessage()
init(k) init(k)
init(m) init(m)
make_request(m, csr=open("req.pem").read()) make_request(m, csr=open(reqfile).read())
sign(open("key.pem").read(), m) sign(open(key_file).read(), m)
print m print m
r=decode(do(m)) r=decode(do(m))
print r print r
@@ -74,7 +83,7 @@ for chall in r.challenge:
print sni_todo print sni_todo
import sni_challenge import sni_challenge
sni_challenge.perform_sni_cert_challenge(sni_todo, "req.pem", "key.pem") sni_challenge.perform_sni_cert_challenge(sni_todo, req_file, key_file)
r=decode(do(k)) r=decode(do(k))
print r print r
@@ -89,8 +98,9 @@ while r.challenge or r.proceed.IsInitialized():
# TODO: there should be a deploy_cert() here. # TODO: there should be a deploy_cert() here.
if r.success.IsInitialized(): if r.success.IsInitialized():
open("cert.pem", "w").write(r.success.certificate) with open(cert_file, "w") as f:
print "Server issued certificate; certificate written to cert.pem" f.write(r.success.certificate)
print "Server issued certificate; certificate written to " + cert_file
elif r.failure.IsInitialized(): elif r.failure.IsInitialized():
print "Server reported failure." print "Server reported failure."
sys.exit(1) sys.exit(1)