import traceback
from binascii import hexlify

#from cancat.j1939 import *
# we can move things into here if we decide this replaces the exiting j1939 modules
import cancat
import struct
from cancat.J1939db import *
from cancat import *
from cancat.vstruct.bitfield import *

import queue
import threading
''' 
This is a J1939 Stack module.
It's purpose is to provide a J1939-capable object which extends (consumes) CanCat's CanInterface module, and provides a J1939 interface.
In the original J1939 module, extended messages were treated like oddities.  This module will work on the premise that all messages (except TP messages) are created equal, and are available in the same queue.  TP messages will be the "specialty" messages which create new, arbitrary sized messages from smaller ones.  If they don't make it through, they don't.  This module is intended to make J1939 attack clients easier to work with, whereas the first module was focused on reverse engineering.  Let's see if we can't come up with something awesome, then possible merge them together in the future.

This module focuses around PGNs.  All messages are handled and sorted by their PS/DA.
'''

J1939MSGS = 1939

PF_RQST =       0xea
PF_TP_DT =      0xeb
PF_TP_CM =      0xec
PF_ADDRCLAIM =  0xee
PF_PROPRIETRY=  0xef
PF_KWP1 =       0xdb
PF_KWP2 =       0xda
PF_KWP3 =       0xce
PF_KWP4 =       0xcd

CM_RTS   =       0x10
CM_CTS   =       0x11
CM_EOM   =       0x13
CM_ABORT =       0xff
CM_BAM   =       0x20

TP_BAM = 20
TP_DIRECT = 10
TP_DIRECT_BROKEN=9

class NAME(VBitField):
    def __init__(self):
        VBitField.__init__(self)
        self.arbaddrcap = v_bits(1)
        self.ind_group = v_bits(3)
        self.vehicle_system_instance = v_bits(4)
        self.vehicle_system = v_bits(7)
        self.reserved = v_bits(1)
        self.function = v_bits(8)
        self.function_instance = v_bits(5)
        self.ecu_instance = v_bits(3)
        self.mfg_code = v_bits(11)
        self.identity_number = v_bits(21)

    def minrepr(self):
        mfgname = mfg_lookup.get(self.mfg_code)
        return "id: 0x%x mfg: %s" % (self.identity_number, mfgname)


def parseName(name):
    namebits= NAME()
    rname = name[::-1]
    namebits.vsParse(rname)
    return namebits

def parseArbid(arbid):
    (prioPlus,
        pf,
        ps,
        sa) = struct.unpack(b'BBBB', struct.pack(b">I", arbid))

    prio = prioPlus >> 2
    edp = (prioPlus >> 1) & 1
    dp = prioPlus & 1

    return prio, edp, dp, pf, ps, sa


def meldExtMsgs(msgs):
    out = []
    length = msgs.get('totsize')
    for arbtup, msg in msgs.get('msgs'):
        out.append(msg[1:])

    outval = b''.join(out)
    outval = outval[:length]

    return outval

### renderers for specific PF numbers
def pf_c9(idx, ts, arbtup, data, j1939):
    b4 = data[3]
    req = "%.2x %.2x %.2x" % ([ord(d) for d in data[:3]])
    usexferpfn = ('', 'Use_Transfer_PGN', 'undef', 'NA')[b4 & 3]
    
    return "Request2: %s %s" % (req,  usexferpgn)

def pf_ea(idx, ts, arbtup, data, j1939):
    (prio, edp, dp, pf, ps, sa) = arbtup
    return "Request: %s" % (hexlify(data[:3]))

# no pf_eb or pf_ec since those are handled at a lower-level in this stack

def pf_ee(idx, ts, arbtup, data, j1939):
    prio, edp, dp, pf, ps, sa = arbtup
    if ps == 255 and sa == 254:
        return 'CANNOT CLAIM ADDRESS'
    
    addrinfo = parseName(data).minrepr()
    return "Address Claim: %s" % addrinfo

def pf_ef(idx, ts, arbtup, data, j1939):
    prio, edp, dp, pf, ps, sa = arbtup
    if dp:
        return 'Proprietary A2'

    return 'Proprietary A1'
    
def pf_ff(idx, ts, arbtup, data, j1939):
    prio, edp, dp, pf, ps, sa = arbtup
    pgn = "%.2x :: %.2x:%.2x - %s" % (sa, pf,ps, hexlify(data))
    return "Proprietary B %s" % pgn

pgn_pfs = {
        0x93:   ("Name Management", None),
        0xc9:   ("Request2",        pf_c9),
        0xca:   ('Transfer',        None),
        0xe8:   ("ACK        ",     None),
        0xea:   ("Request      ",   pf_ea),
        0xeb:   ("TP.DT (WTF?)",    None),
        0xec:   ("TP.CM (WTF?)",    None),
        0xee:   ("Address Claim",   pf_ee),
        0xef:   ("Proprietary",     pf_ef),
        #0xfe:   ("Command Address", None),
        0xff:   ("Proprietary B",   pf_ff),
        }


class J1939Interface(cancat.CanInterface):
    _msg_source_idx = J1939MSGS
    def __init__(self, port=None, baud=cancat.baud, verbose=False, cmdhandlers=None, comment='', load_filename=None, orig_iface=None, process_can_msgs=True, promisc=True):
        self._last_recv_idx = -1
        self._threads = []
        self._j1939_filters = []
        self._j1939_msg_events = {}
        self._j1939queuelock = threading.Lock()
        self._TPmsgParts = {}
        self.maxMsgsPerPGN = 0x200
        self._j1939_msg_listeners = []
        self.promisc = promisc

        cancat.CanInterface.__init__(self, port=port, baud=baud, verbose=verbose, cmdhandlers=cmdhandlers, comment=comment, load_filename=load_filename, orig_iface=orig_iface)

        # setup the message handler event offload thread
        if self._config.get('myIDs') is None:
            self._config['myIDs'] = []

        self._mhe_queue = queue.Queue()
        mhethread = threading.Thread(target=self._mhe_runner)
        mhethread.setDaemon(True)
        mhethread.start()
        self._threads.append(mhethread)

        self.register_handler(CMD_CAN_RECV, self._j1939_can_handler)

        if process_can_msgs:
            self.processCanMessages()

        # restore other config items

    def processCanMessages(self, delete=True):
        for msg in self.recvall(CMD_CAN_RECV):
            self._j1939_can_handler(msg, None)

    def setPromiscuous(self, promisc=True):
        '''
        Determines whether messages not destined for an ID I currently own (self._config['myIDs'] are kept/handled or discarded
        '''
        self.promisc = promisc

    def addID(self, newid):
        if newid not in self._config['myIDs']:
            self._config['myIDs'].append(newid)

    def delID(self, curid):
        if curid in self._config['myIDs']:
            self._config['myIDs'].remove(curid)

    def J1939xmit(self, pf, ps, sa, data, prio=6, edp=0, dp=0):
        if len(data) < 8:
            arbid = emitArbid(prio, edp, dp, pf, ps, sa)
            # print("TX: %x : %r" % (arbid, hexlify(data)))
            self.CANxmit(arbid, data, extflag=1)
            return

        self._j1939xmit_tp(pf, ps, sa, data, prio, edp, dp)

    def _j1939xmit_tp(self, pf, ps, sa, message, prio=6, edp=0, dp=0):
        pgn2 = (edp << 1) | dp
        pgn1 = pf
        if pgn1 < 240:
            pgn0 = 0
        else:
            pgn0 = ps

        msgs = ['%c'%(x+1) + message[x*7:(x*7)+7] for x in range((len(message)+6)//7)]
        if len(msgs) > 255:
            raise Exception("J1939xmit_tp: attempt to send message that's too large")

        cm_msg = struct.pack('<BHBBBBB', CM_RTS, len(message), len(msgs), 0xff, 
                pgn2, pgn1, pgn0)

        arbid = emitArbid(prio, edp, dp, PF_TP_CM, ps, sa)
        # print("TXe: %x : %r" % (arbid, hexlify(cm_msg)))
        self.CANxmit(arbid, cm_msg, extflag=1)
        time.sleep(.01)  # hack: should watch for CM_CTS
        for msg in msgs:
            #self.J1939xmit(PF_TP_DT, ps, sa, msg, prio=prio)
            arbid = emitArbid(prio, edp, dp, PF_TP_DT, ps, sa)
            print("TXe: %x : %r" % (arbid, hexlify(msg)))
            self.CANxmit(arbid, msg, extflag=1)

        # hack: should watch for CM_EOM

    def _reprCanMsg(self, idx, ts, arbtup, data, comment=None):
        #print("_reprCanMsg: %r   %r" % (args, kwargs))

        if comment is None:
            comment = ''

        prio, edp, dp, pf, ps, sa = arbtup

        # give name priority to the Handler, then the manual name (this module), then J1939PGNdb
        pfmeaning, handler = pgn_pfs.get(pf, ('',None))
        nextline = ''

        if handler is not None:
            enhanced = handler(idx, ts, arbtup, data, self)
            if enhanced == cancat.DONT_PRINT_THIS_MESSAGE:
                return enhanced

            if enhanced is not None:
                if type(enhanced) in (list, tuple) and len(enhanced):
                    pfmeaning = enhanced[0]
                    if len(enhanced) > 1:
                        nextline = '\n'.join(list(enhanced[1:]))

                    # if we get multiple lines and the first is DONT_PRINT_THIS_MESSAGE, 
                    # then just return nextline
                    if pfmeaning == cancat.DONT_PRINT_THIS_MESSAGE:
                        return nextline

                    nextline = '\n' + nextline

                else:
                    pfmeaning = enhanced

        elif not len(pfmeaning):
            pgndata = parsePGNData(pf, ps, data)
            pfmeaning = pgndata.get('pgndata').get('Name')
            lines = []

            if self.verbose:
                for spnum, spdict, spunit, spdata, sprepr in pgndata.get('spns'):
                    spnName = spdict.get('Name')
                    lines.append('      SPN(%d): %-20s \t %s' % (spnum, sprepr, spnName))

                if len(lines):
                    nextline = '\n' + '\n'.join(lines)

        return "%.8d %8.3f pri/edp/dp: %d/%d/%d, PG: %.2x %.2x  Source: %.2x  Data: %-18s  %s\t\t%s%s" % \
                (idx, ts, prio, edp, dp, pf, ps, sa, hexlify(data), pfmeaning, comment, nextline)

    def _j1939_can_handler(self, tsmsg, none):
        '''
        this function is run for *Every* received CAN message... and is executed from the 
        XMIT/RECV thread.  it *must* be fast!
        '''
        #print(repr(self), repr(cmd), repr(tsmsg))
        ts, message = tsmsg
        arbid, data = self._splitCanMsg(message)
        arbtup = parseArbid(arbid)
        prio, edp, dp, pf, ps, sa = arbtup

        # if i don't care about this message... bail. (0xef+ is multicast)
        if pf < 0xef and ps not in self._config['myIDs'] and not self.promisc:
            return

        if pf == 0xeb:
            self.queueMessageHandlerEvent(self.eb_handler, arbtup, data, ts)
        elif pf == 0xec:
            self.queueMessageHandlerEvent(self.ec_handler, arbtup, data, ts)
        else:
            self.queueMessageHandlerEvent(self._submitJ1939Message, arbtup, data, ts)

        #print("submitted message: %r" % (hexlify(message)))


    def queueMessageHandlerEvent(self, pfhandler, arbtup, data, ts):
        ''' 
        this is run in the XMIT/RECV thread and is intended to handle offloading the data fast
        '''
        self._mhe_queue.put((pfhandler, arbtup, data, ts))

    def _mhe_runner(self):
        ''' 
        runs the mhe thread, which is offloaded so that the message-handling thread can keep going
        '''
        while not self._config.get('shutdown'):
            while self._config['go']:
                worktup = None
                try:
                    worktup = self._mhe_queue.get(1)
                    if worktup is None:
                        continue

                    pfhandler, arbtup, data, ts = worktup
                    pfhandler(arbtup, data, ts)

                except Exception as e:
                    print("(j1939stack)MsgHandler ERROR: %r (%r)" % (e, worktup))
                    if self.verbose:
                        sys.excepthook(*sys.exc_info())

            time.sleep(1)

    def _submitJ1939Message(self, arbtup, message, timestamp=None):
        '''
        submits a message to the cmd mailbox.  creates mbox if doesn't exist.
        *threadsafe*
        often runs in the MHE thread
        '''
        #print("_submitJ1939Message")
        if timestamp is None:
            timestamp = time.time()

        prio, edp, dp, pf, ps, sa = arbtup
        pgn = (pf<<8) | ps
        datarange = (edp<<1) | dp

        if len(self._j1939_filters):
            lcls = locals()
            for advf in self._j1939_filters:
                try:
                    if not eval(advf, lcls):
                        return
                except Exception as e:
                    print("_submitJ1939Message advfilter ERROR: %r" % e)
                    return

        self._j1939queuelock.acquire()
        try:
            # do we want to break things apart into PGN mboxes at this point?  if so, we have to also allow 
            # subscription at this level for things like sniffing.  Like this:
            handled = False
            for listener in self._j1939_msg_listeners:
                try:
                    contnu = listener(arbtup, message)
                    if contnu: 
                        handled = True

                except Exception as e:
                    self.log('_submitJ1939Message: ERROR: %r' % e)
            
            # check for any J1939 registered handlers (using the default system handlers):
            cmdhandler = self._cmdhandlers.get(J1939MSGS)
            if cmdhandler is not None:
                handled2 = cmdhandler(tsmsg, self)

            if handled and handled2:
                #print("handled")
                return

            ##::: TODO, make this a listener.  if at all...
            #dr = self._messages.get(datarange)
            #if dr is None:
                #dr = {}
                #self._messages[datarange] = dr
            #
            ## factor in multicast vs. unicast...
            #mbox = dr.get(pf)
            #if mbox is None:
                #mbox = []
                #dr[pf] = mbox
                #self._j1939_msg_events[pf] = threading.Event()

            # file in the mailbox
            mbox = self._messages.get(J1939MSGS)