#!/usr/bin/python

'''
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
'''

import argparse
import struct
import sys
import math
import base64
import re
from pylibsrtp import Policy,Session
from collections import namedtuple

__all__ = ['ogg_opus_coder','srtp_hex2hex_coder','ogg_page_coder','RTP_HEADER','RTP_HEADER_FMT','RTP_FIXED_HEADER']
__author__ = "Kamanashis Roy"
__copyright__ = "Copyright (C) 2019 Kamanashis Roy"
__license__ = "GNU Public License version 3"
__version__ = "1.0"

RTP_HEADER = namedtuple('RTP_HEADER', 'FIRST,PAYLOAD_TYPE,SEQUENCE_NUMBER,TIMESTAMP,SSRC')
RTP_HEADER_FMT = '>BBHII'
MAX_RTP_SEQUENCE_NUM = 65535

class hex_coder:
    def __init__(self,verbose):
        self.fd = None
        self.verbose = verbose

    def hexdump(self, content):
        counter = 0
        output = []
        for segment in range((len(content)>>4)+1):
            segment_out = []

            segment_out.append('%06x' % counter)
            for offset in range(0,16):
                pos = (segment<<4)+offset
                if pos >= len(content):
                    break # avoid overflow
                segment_out.append('%02x' % content[pos])
                counter += 1
            output.append(' '.join(segment_out))

        output.append('\n')
        return '\n'.join(output)


    def write_udp(self, udp_payload):
        self.fd.write(self.hexdump(udp_payload))
        self.fd.write("\n")

    def start_file(self, outfile):
        self.fd = open(outfile, 'w')

    def end_file(self):
        self.fd.close()

RTP_FIXED_HEADER = 12
class srtp_hex2hex_coder(hex_coder):
    '''
    Allow record srtp/rtp stream
    '''
    def __init__(self,override_payload_offset=None, rtp_offset=0,srtpkey=None,resrtpkey=None,verbose=False,payload_type=111,filter_ssrc=None):
        #super(ogg_opus_coder, self).__init__(verbose=verbose)
        hex_coder.__init__(self,verbose=verbose)
  
        # setup rtp parameters
        self.override_payload_offset = override_payload_offset
        self.rtp_offset = rtp_offset
        self.payload_type = payload_type

        # setup srtp parameters
        self.srtpkey = srtpkey
        self.resrtpkey = resrtpkey
        self.srtp_session = None
        self.ssrc = None
        self.filter_ssrc = filter_ssrc


    def record_rtp_packet(self, packet,is_last=False):
        if not packet:
            print("No packet")
            return False

        if self.verbose:
            print("Dumping packet ", packet)

        rtp_raw_full = None
        try:
            rtp_raw_full = bytes.fromhex(''.join(packet[self.rtp_offset:]))
        except:
            print("Failed to get hex", packet)
            return False

        if len(rtp_raw_full) < RTP_FIXED_HEADER:
            print("udp payload is too small")
            return False


        if self.verbose:
            print(self.rtp_offset, packet[self.rtp_offset-1], packet[self.rtp_offset], packet[self.rtp_offset+1])

        # decode RTP Header
        rtp_raw_header = rtp_raw_full[:RTP_FIXED_HEADER]

        if self.verbose:
            print(rtp_raw_header.hex())
        rtp = RTP_HEADER._make(struct.unpack(RTP_HEADER_FMT, rtp_raw_header))
        if self.verbose:
            print(rtp)
        if rtp.PAYLOAD_TYPE & 0b10000000:
            # strip marker bit from payload
            rtp = RTP_HEADER(FIRST=rtp.FIRST,PAYLOAD_TYPE=rtp.PAYLOAD_TYPE&0b01111111,SEQUENCE_NUMBER=rtp.SEQUENCE_NUMBER,TIMESTAMP=rtp.TIMESTAMP,SSRC=rtp.SSRC)
            print('stripped marker bit',rtp)

        # Filter the RTP with v=2
        if (rtp.FIRST & 0b11000000) != 0b10000000:
            print("Not an RTP")
            return False

        rtp_exten = rtp.FIRST & 0b10000
        rtp_exten_length = 0
        rtp_csrc  = rtp.FIRST & 0b1111

        # calculate rtp header length
        calc_rtp_header_len = RTP_FIXED_HEADER + rtp_csrc*4
        if rtp_exten:
            exten_start = RTP_FIXED_HEADER+rtp_csrc*4
            exten_raw = rtp_raw_full[exten_start:exten_start+4]
            if len(exten_raw) != 4:
                print("Skipping malformed RTP")
                return False
            rtp_exten_profile,rtp_exten_length = struct.unpack('>HH', exten_raw)
            calc_rtp_header_len += 4 + rtp_exten_length*4

        if self.verbose:
            print("calc_rtp_header_len", calc_rtp_header_len)

        if self.override_payload_offset:
            calc_rtp_header_len = self.override_payload_offset - self.rtp_offset

        # Filter opus
        if self.payload_type and rtp.PAYLOAD_TYPE != self.payload_type:
            print("Skipping payload {rtp_payload_type} while {opus_payload_type} expected.".format(rtp_payload_type=rtp.PAYLOAD_TYPE, opus_payload_type=self.payload_type))
            return False

        if self.filter_ssrc and self.filter_ssrc != rtp.SSRC:
            print("Skipping ssrc={rtp_ssrc} while {filter_ssrc} expected".format(rtp_ssrc=rtp.SSRC,filter_ssrc=self.filter_ssrc))
            return False

        if len(rtp_raw_full) < (calc_rtp_header_len+1):
            print("Empty payload")
            return False

        if self.srtpkey:
            if self.ssrc != rtp.SSRC:
                #if not self.srtp_session:
                #    self.srtp_session = Session()
                self.srtp_session = Session()
                print("using key [%s]" % self.srtpkey)
                srtpkey = base64.b64decode(self.srtpkey)
                plc = Policy(key=srtpkey,ssrc_value=rtp.SSRC,ssrc_type=Policy.SSRC_ANY_INBOUND)
                print(plc)
                self.srtp_session.add_stream(plc)
                self.ssrc = rtp.SSRC
            try:
                rtp_raw_full = self.srtp_session.unprotect(rtp_raw_full)
            except:
                print("decrypt fail seq={sequence}, ssrc={ssrc}".format(sequence=rtp.SEQUENCE_NUMBER,ssrc=rtp.SSRC))
                '''
                if self.resrtpkey:
                    srtpkey = base64.b64decode(self.resrtpkey)
                    plc = Policy(key=srtpkey,ssrc_value=rtp.SSRC,ssrc_type=Policy.SSRC_ANY_INBOUND)
                    self.srtp_session.add_stream(plc)
                    print("Using restrpkey here from next packet")
                '''
                return False


        self.write_udp(bytes.fromhex(''.join(packet[:self.rtp_offset])) + rtp_raw_full)
        return True


    def record_rtp_file(self, infile, outfile, hexstring=False):
        '''
        Convert an RTP hexdump into a recorded file
        '''
        self.start_file(outfile)

        if hexstring:
            with open(infile, 'r') as rtp_fd:
                packet_counter = 0
                success_counter = 0
                for xline in rtp_fd:
                    xline = xline.strip()
                    packet = [xline[idx:idx+2] for idx in range(0,len(xline),+2)]
                    packet_counter += 1
                    if self.record_rtp_packet(packet):
                        success_counter += 1
        else:
            re_valid_hex = re.compile(r'^[0-9A-Fa-f]{1,8}\s{1,4}[0-9A-Fa-f]{2}')
            with open(infile, 'r') as rtp_fd:
                packet_counter = 0
                success_counter = 0
                packet = []
                for xline in rtp_fd:
                    if not xline or not re_valid_hex.match(xline):
                        if packet:
                            packet_counter += 1
                            if self.record_rtp_packet(packet):
                                success_counter += 1
                            packet = []
                    else:
                        content = xline.split()
                        #print(len(content))
                        content.pop(0) # skip the segment column
                        packet.extend(content)

                if packet:
                    packet_counter += 1
                    if self.record_rtp_packet(packet,is_last=True):
                        success_counter += 1

                print("Written %d out of %d packets" % (success_counter, packet_counter))
        self.end_file()

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='RTP hex-dump reader/writer tool (useful to decrypt rtp)')
    parser.add_argument('-v', '--verbose'
        , required=False
        , action='store_true'
        , help='Show verbose output')

    record_group = parser.add_argument_group('Record rtp/srtp stream')
    record_group.add_argument('-r', '--recordfile'
        , required=False
        , help='hex output file')

    record_group.add_argument('--payloadoffset'
        , required=False
        , type=int
        , default=None
        , help='Forced payload offset, should be 12+rtpoffset when no extension. Currently extension is calculated automatically in case the offset not forced., see also --rtpoffset')

    record_group.add_argument('--rtpoffset'
        , required=False
        , type=int
        , default=0
        , help='Offset of rtp header')

    record_group.add_argument('--ssrc'
        , required=False
        , type=int
        , default=None
        , help='Specify ssrc as it appears in SDP. It is used to filter-out specific stream')
 
                   
    playfile_group = parser.add_argument_group('Playback opus file into hex-stream')
    playfile_group.add_argument('--playfile'
        , required=False
        , help='ogg formatted opus input file that is streamed into rtp hex file')

    playfile_group.add_argument('--streamseq'
        , required=False
        , type=int
        , default=0
        , help='Specify starting sequence while creating playback rtp stream')
 
    # Support encryption/decryption
    parser.add_argument('-k', '--srtpkey'
        , required=False
        , default=None
        , help='srtpkey to encrypt/decrypt the rtp')

    parser.add_argument('--resrtpkey'
        , required=False
        , default=None
        , help='srtp re-key/second-key to decrypt the rtp\n'
            'It is used when the srtpkey fails to decrypt valid rtp-payload.')
                   
    parser.add_argument('--payloadtype'
        , required=False
        , type=int
        , default=111
        , help='Specify payload type as it appears in SDP. It is used to filter-out RTCP packets')
                   

    parser.add_argument('-x', '--hexfile'
        , required=False
        , help='Text output of pcap, hint `tshark -x -r file.pcap -Y "rtp && udp.srcport == myport" | cut -d " " -f 1-20 > strm_from_myport.txt`\n'
            'In case of recording, it is input file. In case of playback, it is the output file')

    parser.add_argument('--hexstring'
        , action='store_true'
        , help='Text output of pcap, hint `tshark -x -r file.pcap -Y "rtp && udp.srcport == myport" -T fields -e data > strm_from_myport.txt`\n'
            'In case of recording, it is input file. Playback is not supported in this format.')



    args = parser.parse_args()
    if args.verbose:
        print(args)
    if args.hexfile and args.recordfile:
        hextohex = srtp_hex2hex_coder(override_payload_offset=args.payloadoffset,rtp_offset=args.rtpoffset,verbose=args.verbose,srtpkey=args.srtpkey,resrtpkey=args.resrtpkey,payload_type=args.payloadtype,filter_ssrc=args.ssrc)
        hextohex.record_rtp_file(args.hexfile, args.recordfile, args.hexstring)


