diff --git a/lib/modules/blockchain_process/proxy.js b/lib/modules/blockchain_process/proxy.js
index a16be560..02b070af 100644
--- a/lib/modules/blockchain_process/proxy.js
+++ b/lib/modules/blockchain_process/proxy.js
@@ -36,7 +36,7 @@ const parseResponse = function(ipc, resBody){
     } catch(e) {
         return; // Response is not a json. Do nothing
     }
-
+    console.log(jsonO);
     if(commList[jsonO.id]){
         commList[jsonO.id].transactionHash = jsonO.result;
         transactions[jsonO.result] = {commListId: jsonO.id};
diff --git a/lib/modules/coverage/contract_source.js b/lib/modules/coverage/contract_source.js
index 29bb1033..9b65e64c 100644
--- a/lib/modules/coverage/contract_source.js
+++ b/lib/modules/coverage/contract_source.js
@@ -95,6 +95,7 @@ class ContractSource {
       throw new Error('Error generating coverage: solc output was not assigned');
 
     var coverage = {
+      code: this.body.split("\n"),
       l: {},
       path: this.file,
       s: {},
@@ -130,24 +131,31 @@ class ContractSource {
 
           var location = this.sourceMapToLocations(node.src);
           var trueBranchLocation = this.sourceMapToLocations(node.trueBody.src);
-          var falseBranchLocation = this.sourceMapToLocations(node.falseBody.src);
+
+          var falseBranchLocation;
+          if(node.falseBody) {
+            falseBranchLocation = this.sourceMapToLocations(node.falseBody.src);
+          } else {
+            falseBranchLocation = trueBranchLocation;
+          }
 
           coverage.branchMap[node.id] = {
             type: 'if',
             locations: [trueBranchLocation, falseBranchLocation],
-            line: location.start.line,
-          }
+            line: location.start.line
+          };
 
           children = [node.condition]
-            .concat(node.trueBody.statements)
-            .concat(node.falseBody.statements);
+            .concat(node.trueBody.statements);
+
+          if(node.falseBody) children = children.concat(node.falseBody.statements);
 
           markLocations = [location, trueBranchLocation, falseBranchLocation];
 
           if(node.trueBody.statements[0])
             node.trueBody.statements[0]._parent = {type: 'b', id: node.id, idx: 0}
 
-          if(node.falseBody.statements[0])
+          if(node.falseBody && node.falseBody.statements[0])
             node.falseBody.statements[0]._parent = {type: 'b', id: node.id, idx: 1}
 
           sourceMapToNodeType[node.src] = [{
@@ -225,7 +233,7 @@ class ContractSource {
             body: coverage.fnMap[node.id],
           }];
 
-          children = node.body.statements;
+          if(node.body) children = node.body.statements;
           markLocations = [location];
           break;
 
@@ -244,11 +252,26 @@ class ContractSource {
 
     } while(nodesRequiringVisiting.length > 0);
 
+    var contractMatches = true;
     for(var contractName in this.contractBytecode) {
       var bytecode = this.contractBytecode[contractName];
+
+      // Try to match the contract to the bytecode. If it doesn't,
+      // then we bail.
+      trace.structLogs.forEach((step) => {
+        if(!contractMatches) return;
+
+        if(!bytecode[step.pc]) {
+          // This contract won't match so we bail
+          contractMatches = false;
+          return;
+        }
+      })
+
+      if(!contractMatches) break;
+
       trace.structLogs.forEach((step) => {
         var step = bytecode[step.pc];
-        step.seen = true;
 
         if(!step.sourceMap || step.sourceMap == '')
           return;
diff --git a/lib/modules/coverage/index.js b/lib/modules/coverage/index.js
index 21a3a1dc..f71c4917 100644
--- a/lib/modules/coverage/index.js
+++ b/lib/modules/coverage/index.js
@@ -1,5 +1,15 @@
+/*global web3*/
+const fs = require('fs');
+const process = require('process');
+
 const ContractSources = require('./contract_sources');
 
+// Set up the web3 extension
+web3.extend({
+  property: 'debug',
+  methods: [{name: 'traceTransaction', call: 'debug_traceTransaction', params: 2}]
+});
+
 class CodeCoverage {
   constructor(embark, _options) {
     this.events = embark.events;
@@ -8,9 +18,14 @@ class CodeCoverage {
     embark.events.on('contracts:compile:solc', this.compileSolc.bind(this));
     embark.events.on('contracts:compiled:solc', this.compiledSolc.bind(this));
     embark.events.on('contracts:run:solc', this.runSolc.bind(this));
-    embark.registerActionForEvent('contracts:deploy:afterAll', this.deployed.bind(this));
-
     embark.events.on('block:header', this.runSolc.bind(this));
+
+    this.seenTransactions = {};
+
+    var self = this;
+    process.on('exit', (code) => {
+      fs.writeFileSync('/tmp/coverage.json', JSON.stringify(this.coverageReport));
+    });
   }
 
   compileSolc(input) {
@@ -27,18 +42,18 @@ class CodeCoverage {
     this.contractSources.parseSolcOutput(output);
   }
 
-  runSolc(receipt) {
-    console.log('runSolc');
-    console.dir(receipt);
-    //this.contractSources.generateCodeCoverage(trace);
-  }
+  async runSolc(receipt) {
+    var block = await web3.eth.getBlock(receipt.number);
+    for(var i in block.transactions) {
+      var txHash = block.transactions[i];
 
-  deployed(cb) {
-    this.events.request('contracts:list', (error, contracts) => {
-      console.dir(error);
-      console.dir(contracts);
-      cb();
-    });
+      if(this.seenTransactions[txHash]) return;
+
+      this.seenTransactions[txHash] = true;
+      var trace = await web3.debug.traceTransaction(txHash, {});
+      var cov = this.contractSources.generateCodeCoverage(trace);
+      this.coverageReport = cov;
+    }
   }
 }
 
diff --git a/test/coverage.js b/test/coverage.js
index 12af9609..555a5d8b 100644
--- a/test/coverage.js
+++ b/test/coverage.js
@@ -1,4 +1,5 @@
-const {assert, expect} = require('chai');
+/*global describe, it*/
+const {assert} = require('chai');
 const fs = require('fs');
 const path = require('path');
 const sinon = require('sinon');
@@ -39,7 +40,7 @@ describe('ContractSources', () => {
 
     it('should throw an error when the file does not exist', (done) => {
       assert.throws(() => {
-        new ContractSources(['fixtures/404.sol'])
+        new ContractSources(['fixtures/404.sol']);
       }, 'Error loading fixtures/404.sol: ENOENT');
 
       done();