import dshell
import util
import binascii
import json
from struct import unpack

# Magic values used to determine Bitcoin Network Type
# Bitcoin Testnet is an alternative blockchain used for testing
MAGIC_VALS = {'F9 BE B4 D9': 'BITCOIN-MAIN',
              'FA BF B5 DA': 'BITCOIN-TESTNET',
              '0B 11 09 07': 'BITCOIN-TESTNET3'}


class DshellDecoder(dshell.TCPDecoder):

    def __init__(self):
        dshell.TCPDecoder.__init__(self,
                            name='bitcoin',
                            description='Extract Bitcoin traffic, including Stratum mining protocol (pooled) traffic',
                            longdescription='''
The bitcoin decoder will extract any Bitcoin traffic attempting to find and output: 
    Client/server IP addresses, src/dst port numbers, MAC addresses of the machines
    used in the Bitcoin communication/transactions, timestamps of the packets,
    packet payload sizes in KB, and the Network type 
    ('Bitcoin Main' if Bitcoin data traffic)  

Additionally for Stratum mining, it will attempt to extract:
    Bitcoin miner being used, transaction methods used in each connection 
    (mining.notify, mining.authorize, mining.get_transaction, mining.submit, etc.), 
    User ID (Auth ID) used to access the Bitcoin mining pool, and possibly the password 
    used to connect to the pool if it is stored in the JSON of the payload.

    Also, the: 
    range of job IDs (per connection), previous block hash, generation transaction (part 1), 
    generation transaction (part 2), merkle tree branches (hashes), block version, 
    and the hash difficulty (n-bits)
    (The generation transactions and merkle tree branches are only optionally outputted 
     to a file: See Example (3) below)
 
    Note (1): The first time that all of this Stratum mining information is 
              collected (per connection), all of the packets decoded after this 
              point from within the same connection (same exact sip, dip, sport, dport) 
              will continue to output the same collection of information since it 
              will be the same.

    Note (2): The gen_tx1 and gen_tx2 fields enable the miner to build the coinbase
              transaction for the block by concatentating gen_tx1, the extranonce1 
              at the start of gen_tx1, the extranonce2 generated by the miner, and 
              gen_tx2 (hashes with scriptPubKeys)

    Note (3): Some pools recommend miners use their Bitcoin wallet ID 
              (address used for payment) as their 'Auth ID'. This will be easily 
              spotted as it is an address consisting of 26-35 alphanumeric characters, 
              and it always begins with either the number '1' or '3'


    For additional information: 
        Bitcoin Protocol:        <https://en.bitcoin.it/wiki/Protocol_documentation>
        Stratum Mining Protocol: <https://en.bitcoin.it/wiki/Stratum_mining_protocol>


--------Main ports / Some secondary ports used for Bitcoin Traffic---------
    Bitcoin Main traffic uses port 8333
    Bitcoin Testnet uses port 18333
    Several pools use ports 3333, 8332, 8337
    The other ports checked are known ports used by specific BTC mining pools
    Other Bitcoin pools utilize alternate ports (even 80 / 443)


Examples:
    
    (1) Basic usage:

        decode -d bitcoin <pcap>

    (2) If pcap starts in middle of connection (then we need to ignore TCP handshake):
        
        decode -d bitcoin <pcap> --bitcoin_ignore_handshake

    (3) Saving Generation Transaction Data and Merkle Branches to a specified file:

        decode -d bitcoin <pcap> --bitcoin_gentx='foo.txt.'
''',
                            filter='''tcp and port (3332 or 3333 or 3334 or 3335 or
                                                    4008 or 4012 or 4016 or 4024 or
                                                    4032 or 4048 or 4064 or 4096 or
                                                    4128 or 4256 or 5050 or 7033 or
                                                    7065 or 7129 or 7777 or 8332 or
                                                    8333 or 8334 or 8336 or 8337 or 
                                                    8344 or 8347 or 8361 or 8888 or 
                                                    9332 or 9337 or 9999 or 11111 or 
                                                    12222 or 17777 or 18333)''',
                            author='dek',
                            optiondict={
                                    'gentx' : {'type' : 'str', 'default' : '', 'help' : 'The name of the file to output the fields used to generate the block transaction (gen_tx1, gen_tx2, merkle_branches).'},
                            }
                            )
        self.auth_ids = {}
        self.notify_params = {}
        self.methods = {}
        self.miners = {} 
        self.job_ids = {}
        self.smac = None
        self.dmac = None
        self.bc_net = None
        self.size = 0
        self.JSON = False
        self.NOTIFY = False

        
    # blobHandler to reassemble the packets in the traffic
    # Bitcoin traffic uses TCP
    def blobHandler(self, conn, blob):

        try:
            data = blob.data()
        except:
            self.warn('could not parse session data')
            return
        
        # Only continue if the packet contains data 
        if not data:
            return

        # Default mining.notify fields to None
        job_id = None 
        prev_blk_hash = None 
        gen_tx1 = None  
        gen_tx2 = None 
        merkle_branches = None
        blk_ver = None 
        difficulty = None 
        curr_time = None 
        clean_jobs = None 
        
        # If the payload contains JSON
        if data.startswith('{"'):
            self.JSON = True
            try:
                # split JSON objects by newline
                for rawjs in data.split("\n"):
                    if rawjs:
                        js = json.loads(rawjs)
                        try:
                            if "method" in js and js["method"]:
                                # Create a dictionary of sets of mining methods 
                                #   indexed by their associated conn.addr (sip, dip, sport, dport) 
                                self.methods.setdefault(conn.addr, set([])).add(js["method"])

                                if js["method"] == "mining.subscribe":
                                    self.miners[conn.addr] = js["params"][0]
                            
                                if js["method"] == "mining.authorize":
                                    if "params" in js and js['params'][0]:
                                        # Grab the Bitcoin User ID (sometimes a wallet id) 
                                        #   which is being authorized
                                        self.auth_ids[conn.addr] = js["params"][0]

                                        if js['params'][1]:
                                            self.auth_ids[conn.addr] = "".join(( 
                                                self.auth_ids[conn.addr], " / ",  str(js['params'][1]) ))
                                            
                                if js["method"] == "mining.notify":
                                    self.NOTIFY = True
                                    if "params" in js and js['params']:
                                        job_id, prev_blk_hash, gen_tx1, gen_tx2, merkle_branches, blk_ver, difficulty, curr_time, clean_jobs = js['params']
                                        self.job_ids.setdefault(conn.addr, []).append(job_id)
                                        self.notify_params[conn.addr] = [self.job_ids, prev_blk_hash, gen_tx1,
                                                                         gen_tx2, merkle_branches,
                                                                         blk_ver, difficulty, curr_time,
                                                                         clean_jobs]

                        except KeyError as e:
                            self.warn("{} - Error extracting auth ID".format(str(e)))
            except ValueError as e:
                self.warn('{} - json data not found'.format(str(e)))
                return
         
        # Grab the first 4 bytes of the payload to search for the magic values
        #   used to determine which Bitcoin network is being accessed
        # Additionally, reformat bytes
        try:
            magic_val = binascii.hexlify(data[0:4]).upper()
            magic_val = ' '.join([magic_val[i:i+2] for i in range(0, len(magic_val), 2)])
        except:
            self.warn('could not parse session data')
            return 

        # Attempt to translate first 4 bytes of payload 
        #   into a Bitcoin (bc) network type
        try:
            self.bc_net = str(MAGIC_VALS[magic_val])  
        except:
            self.bc_net = 'N/A'


        # Pull pertinent information from packet's contents
        self.size = '{0:.2f}'.format((len(blob.data())/1024.0))
        self.smac = conn.smac
        self.dmac = conn.dmac


        # Truncate the list Job IDs per connection for printing purposes if JSON 
        #   data was found in the blob
        if self.JSON and self.NOTIFY:
            jids_end = (len(self.notify_params[conn.addr][0][conn.addr]) - 1)
            if jids_end >= 1:
                self.notify_params[conn.addr][0][conn.addr] = [ self.notify_params[conn.addr][0][conn.addr][0], 
                                                                '...',
                                                                self.notify_params[conn.addr][0][conn.addr][jids_end] ]
            elif jids_end >= 0:
                self.notify_params[conn.addr][0][conn.addr] = [self.notify_params[conn.addr][0][conn.addr][0]] 

            # Reset the JSON data found in the current blob boolean
            #   and the mining.notify method type found in payload boolean 
            self.JSON = False
            self.NOTIFY = False

        # If able to pull the Bitcoin Pool User ID (sometimes a wallet ID)
        # Also if the transcation is mining.notify (seen in Stratum mining)
        #   then output the current Block information
        if (self.size and self.smac and self.dmac and 
                self.miners.get(conn.addr, None) and self.methods.get(conn.addr, None) and self.auth_ids.get(conn.addr, None)
                and self.notify_params.get(conn.addr, None) and not self.gentx): 
            self.alert("\n\tSRC_MAC: \t{0:<20} \n\tDST_MAC: \t{1:<20} \n\tSIZE: \t\t{2:>3}KB\n\t" 
                       "MINER: \t\t{3:<20} \n\tMETHODS: \t{4:<25} \n\tUSER ID/PW: \t{5:<50}\n\t"
                       "JOB IDs: \t{6:<20} \n\tPREV BLK HASH: \t{7:<65} \n\tBLOCK VER: \t{8:<15}\n\t"
                       "HASH DIFF: \t{9:<10}\n\n".format(
                self.smac, self.dmac, self.size, 
                self.miners[conn.addr], self.methods[conn.addr], self.auth_ids[conn.addr],
                self.notify_params[conn.addr][0][conn.addr], self.notify_params[conn.addr][1], 
                self.notify_params[conn.addr][5], self.notify_params[conn.addr][6]), 
                ts=blob.starttime, sip=conn.sip, dip=conn.dip, sport=conn.sport, 
                dport=conn.dport, direction=blob.direction)


        # If able to pull the Bitcoin Pool User ID (sometimes a wallet ID)
        # Also if the transcation is mining.notify (seen in Stratum mining) and the user
        #   specifies that they want to save the fields used to generate the block transaction 
        #   (gen_tx1, gen_tx2 (hashes with scriptPubKeys), merkle tree branches), then
        # output all information possible, and write the gentx information to the specified file
        elif (self.size and self.smac and self.dmac and 
                self.miners.get(conn.addr, None) and self.methods.get(conn.addr, None) and self.auth_ids.get(conn.addr, None)
                and self.notify_params.get(conn.addr, None) and self.gentx): 
            self.alert("\n\tSRC_MAC: \t{0:<20} \n\tDST_MAC: \t{1:<20} \n\tSIZE: \t\t{2:>3}KB\n\t" 
                       "MINER: \t\t{3:<20} \n\tMETHODS: \t{4:<25} \n\tUSER ID/PW: \t{5:<50}\n\t"
                       "JOB IDs: \t{6:<20} \n\tPREV BLK HASH: \t{7:<65} \n\tBLOCK VER: \t{8:<15}\n\t"
                       "HASH DIFF: \t{9:<10}\n\n".format(
                self.smac, self.dmac, self.size, 
                self.miners[conn.addr], self.methods[conn.addr], self.auth_ids[conn.addr],
                self.notify_params[conn.addr][0][conn.addr], self.notify_params[conn.addr][1], 
                self.notify_params[conn.addr][5], self.notify_params[conn.addr][6]), 
                ts=blob.starttime, sip=conn.sip, dip=conn.dip, sport=conn.sport, 
                dport=conn.dport, direction=blob.direction)

            # Write the verbose block information (gen tx1/2, merkle branches) gathered 
            #   from mining.notify payloads to the command-line specified output file
            # The extra information (JOB ID, BLOCK VER, etc.) will be useful in matching the 
            #   information outputted by the alerts to the payload containing the 
            #   generation transaction info and merkle branches
            fout = open(self.gentx, "a+") 
            fout.write(("\nJOB IDs: \t\t{0:<20} \nPREV BLK HASH: \t{1:<65} \n\nGEN TX1: \t\t{2:<20}" 
                        "\n\nGEN TX2: \t\t{3:<20} \n\nMERKLE BRANCHES: {4:<20} \n\nBLOCK VER: \t\t{5:<20}" 
                        "\nHASH DIFF: \t\t{6:<10}\n").format(
                        self.notify_params[conn.addr][0][conn.addr], self.notify_params[conn.addr][1],
                        self.notify_params[conn.addr][2], self.notify_params[conn.addr][3],
                        self.notify_params[conn.addr][4], self.notify_params[conn.addr][5],
                        self.notify_params[conn.addr][6]))
            fout.write(("\n" + "-"*100)*2)


        # Else if we dont have Bitcoin User IDs, or Block information 
        #   and the user doesn't want verbose block information (gentx)
        elif (self.size and self.smac and self.dmac and self.bc_net): 
            self.alert("\n\tNETWORK: \t{0:<15} \n\tSRC_MAC: \t{1:<20} \n\tDST_MAC: \t{2:<20} \n\tSIZE: \t\t{3:>3}KB\n\n".format(
            self.bc_net, self.smac, self.dmac, self.size), ts=blob.starttime, sip=conn.sip, 
            dip=conn.dip, sport=conn.sport, dport=conn.dport, direction=blob.direction)
        

if __name__ == '__main__':
    dObj = DshellDecoder()
    print dObj
else:
    dObj = DshellDecoder()
