#!/usr/bin/python

import socket
import re
import bitstring
import json
import hashlib
from xml.dom.minidom import parseString

METADATA_FORMATS = ['JSON', 'onvif']

SERVER_IP = '192.168.0.19'  # VCAcore IP ADDRESS
PORT = 8554  # VCAcore RTSP PORT
CHANNEL_ID = 2  # VCAcore CHANNEL ID
USERNAME = 'admin'  # VCAcore USERNAME AND PASSWORD
PASSWORD = 'admin'
FORMAT = METADATA_FORMATS[1] # DESIRED METADATA FORMAT
CLIENT_PORTS = [8000, 8001]  # CLIENT PORTS TO BE USED FOR RECEIVING DATA


def send_rtsp_method(socket_connection, method_string, address):
    """Sends a METHOD string to the RTSP Server. Handles VCAcore's Digest Authentication"""
    print(f"*** SENDING {method_string.split()[0]} ***")
    socket_connection.send(str.encode(method_string))
    received_data = socket_connection.recv(4096)
    print("*** GOT ****")
    # CHECK FOR AUTHENTICATION. IF FOUND, CALCULATE RESPONSE HASH AND INJECT authentication INTO THE DESCRIBE METHOD
    auth, realm, nonce = get_authentication(received_data)
    print_rec(received_data)
    if nonce is not None:
        ha1 = hashlib.md5((USERNAME + ":" + realm + ":" + PASSWORD).encode('utf-8')).hexdigest()
        ha2 = hashlib.md5((method_string.split()[0] + ":" + address).encode('utf-8')).hexdigest()
        response = hashlib.md5((ha1 + ":" + nonce + ":" + ha2).encode('utf-8')).hexdigest()
        authentication = "Authorization: " + auth + " username=\"" + USERNAME + "\", realm=\"" + realm + "\", nonce=\"" \
                         + nonce + "\", uri=\"" + address + "\", response=\"" + response + "\"\r\n"
        # INJECT AUTHENTICATION INTO THE method_string
        new_method_string = method_string.replace('User-Agent', authentication + 'User-Agent')
        socket_connection.send(str.encode(new_method_string))
        received_data = socket_connection.recv(4096)
        print_rec(received_data)
    return received_data


def get_authentication(data):
    """Searching for auth type, realm and nonce info from the rtsp strings using regular expressions"""
    auth_pat = re.compile("WWW-Authenticate: .* realm=")
    realm_pat = re.compile("realm=\".*\",")
    nonce_pat = re.compile("nonce=\".*\"")
    auth_results = auth_pat.findall(data.decode())
    realm_results = realm_pat.findall(data.decode())
    nonce_results = nonce_pat.findall(data.decode())
    if len(realm_results) > 0 and len(nonce_results) > 0:
        return auth_results[0][18:-7], realm_results[0][7:-2], nonce_results[0][7:-1]
    else:
        # print('NO REALM OR NONCE FOUND')
        return None, None, None


def print_rec(data):
    """Pretty-printing rtsp strings"""
    recs = data.decode().split('\r\n')
    for rec in recs:
        print(rec)


def get_session_id(data):
    """Search session id from rtsp strings"""
    recs = data.decode().split('\r\n')
    for rec in recs:
        ss = rec.split()
        if ss[0].strip() == "Session:":
            return ss[1].split(";")[0].strip()


def get_stream(data, metadata_format):
    """Search stream id for JSON metadata from rtsp strings"""
    if metadata_format not in METADATA_FORMATS:
        print(f"{metadata_format} format not supported")
        return ""
    recs = data.decode().split('\r\n')
    # PARSE HEADERS
    streams = []
    stream = None
    # LOOP THROUGH AND STRIP OUT THE STREAMS (MEDIA NAME AND TRANSPORT ADDRESS)
    for rec in recs:
        if rec[0:2] == "m=":
            if stream:
                streams.append(stream) 
            stream = {}
            header_name, header_value = rec.split("=", 1)
            stream[header_name.strip()] = header_value.strip()
        if ":" in rec:
            if stream:
                header_name, header_value = rec.split(":", 1)
                stream[header_name.strip()] = header_value.strip()
    # ADD LAST STEAM TO STREAMS
    streams.append(stream)
    # FIND THE JSON STREAM IN THE STREAMS
    for stream in streams:
        if metadata_format in stream['a=rtpmap']:
            return stream['a=control']


def digest_metadata(data):
    """This routine takes a UDP packet, i.e. a string of bytes and ..
    (a) strips off the RTP header
    (b) returns the metadata (in bytes) and the marker indicating if this is a new metadata payload"""
    # PROCESS RTP PACKET HEADER
    bt = bitstring.BitArray(bytes=data)
    version = bt[0:2].uint      # RTP PROTOCOL VERSION
    p = bt[2]                   # PADDING
    x = bt[3]                   # EXTENSION
    cc = bt[4:8].uint           # CSRC COUNT
    m = bt[8]                   # MARKER
    pt = bt[9:16].uint          # PAYLOAD TYPE
    sn = bt[16:32].uint         # SEQUENCE NUMBER
    timestamp = bt[32:64].uint  # TIMESTAMP
    ssrc = bt[64:96].uint       # ssrc IDENTIFIER

    lc = 12  # BYTE INDEX SET AT THE FIXED OFFSET
    bc = 12 * 8  # BIT INDEX
    cids = []
    for j in range(cc):
        cids.append(bt[bc:bc + 32].uint)
        lc = lc + 4
        bc = bc + 32

    print(f"Version: {version}, Padding: {p}, Header Extension: {x}, CSRC Count: {cc}, Marker: {m}, Payload Type: {pt} \
          Sequence Num: {sn}, Timestamp: {timestamp}, Sync Source Id: {ssrc}, CSRC ids: {cids}")
    if x:  # HEADER EXTENSION
        hid = bt[bc:bc + 16].uint
        lc = lc + 2
        bc = bc + 16

        hlen = bt[bc:bc + 16].uint
        lc = lc + 2
        bc = bc + 16
        print(f"ext. header id: {hid}, header len: {hlen}")

        hst = bt[bc:bc + 32 * hlen]
        lc = lc + (4 * hlen)
        bc = bc + (32 * hlen)
    # RETURN JUST THE JSON DATA AND IF THIS IS THE FIRST PACKET
    return data[lc:], m


if __name__ == "__main__":
    adr = "rtsp://" + SERVER_IP + ":" + str(PORT) + "/channels/" + str(CHANNEL_ID)

    # CREATE A TCP SOCKET FOR CONNECTING TO THE RTSP STREAM
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((SERVER_IP, PORT))

    options = "OPTIONS " + adr + " RTSP/1.0\r\nCSeq: 1\r\nUser-Agent: python\r\nAccept: application/sdp\r\n\r\n"
    decoded_data = send_rtsp_method(s, options, adr)

    describe = "DESCRIBE " + adr + " RTSP/1.0\r\nCSeq: 2\r\nUser-Agent: python\r\nAccept: application/sdp\r\n\r\n"
    received_string = send_rtsp_method(s, describe, adr)
    stream = get_stream(received_string, FORMAT)

    setup = "SETUP " + adr + "/" + str(stream) + \
            " RTSP/1.0\r\nCSeq: 3\r\nUser-Agent: python\r\nTransport: RTP/AVP;unicast;client_port=" \
            + str(CLIENT_PORTS[0]) + "-" + str(CLIENT_PORTS[1]) + "\r\n\r\n"

    received_string = send_rtsp_method(s, setup, adr)

    # CREATE CONNECTION TO THE DESIRED METADATA STREAM
    s1 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s1.bind(("", CLIENT_PORTS[0]))
    s1.settimeout(5)  # TIMEOUT IF DEAD FOR 5 SECONDS

    # PORT IS OPEN TO RECEIVE json DATA, SEND PLAY OPTION
    play = "PLAY " + adr + " RTSP/1.0\r\nCSeq: 5\r\nUser-Agent: python\r\nSession: SESID\r\nRange: npt=0.000-\r\n\r\n"
    received_string = send_rtsp_method(s, play.replace("SESID", str(get_session_id(received_string))), adr)

    print("** GET AND DECODING DATA **")
    data = None
    while True:
        received_string = s1.recv(1500)
        print("read", len(received_string), "bytes")
        current_data, marker = digest_metadata(received_string)
        # IF FIRST PACKET TO BE RECEIVED
        if data is None:
            # SET THE DATA
            data = current_data
        else:
            data = data + current_data

        if marker:
            # MAKING THE CURRENT data COMPLETE AND READY TO BE PROCESSED
            if FORMAT == METADATA_FORMATS[0]:
                valid_json = json.loads(data.decode())
                print(valid_json)
            elif FORMAT == METADATA_FORMATS[1]:
                data = data.decode()
                valid_xml = parseString(data)
                print(valid_xml.toprettyxml())
            # RESET
            data = None
    # BEFORE CLOSING SOCKETS, TEARDOWN SHOULD BE SENT VIA RTSP
    s.close()
    s1.close()