/* rtpstream.c  ver. 3.0  7. Apr. 2015
 * This utility converts raw UDP MPEG Transport Stream from legacy encoders
 * into RFC-compliant stream with RTP headers suitable for transmission
 * over IP networks. */



#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/time.h>

typedef signed long long    s64;

    
void usage()
{
    fprintf(stderr,
        "Usage: rtpstream -i inaddr:port -o outaddr:port [-s srcaddr[:port]] -t ttl\n"
        "    inaddr|outaddr are typically multicast IP addresses\n"
        "    use -s to accept packets only from the specified source\n");
}

parsearg(arg,addr,port)
    char *arg;
    in_addr_t *addr;
    ushort *port;
{
    char buf[21];
    int pnum=0;

    if (!sscanf(arg, "%20[0-9.]:%d", &buf, &pnum))
        return -1;
    *addr = inet_addr(buf);
    *port = htons(pnum);
    return 0;
}

main(argc,argv)

    int argc;
    char *argv[];
{
    struct sockaddr_in inaddr, outaddr, srcaddr;
    int c, insock, outsock, ttl=0, len, srclen; 
    in_addr_t opt, inip=0, outip=0, srcip=0;
    ushort inport=0, outport=0, srcport=0;
    uint8_t buf[1550];
    uint16_t rtpseq=0;
    s64 rtpts;
    struct timeval tv;

    while ((c = getopt(argc, argv, "i:o:s:t:")) != -1)
        switch (c) {
        case 'i':
            if (!parsearg(optarg, &inip, &inport)) break;
        case 'o':
            if (!parsearg(optarg, &outip, &outport)) break;
        case 's':
            if (!parsearg(optarg, &srcip, &srcport)) break;
        case 't':
            ttl=atoi(optarg);
            break;
        default:
            usage();
            return -1;
        }
      
    if (!inip || !inport || !outip || !outport | !ttl)
    {
        usage();
        return -1;
    }

    if ((insock=socket(AF_INET,SOCK_DGRAM,0)) < 0)
    {
        perror("can't create input socket");
        return -1;
    }

    /* We may want to reuse an already used socket */

    opt = 1;
    setsockopt( insock, SOL_SOCKET, SO_REUSEADDR, (void *) &opt,
                    sizeof( opt ) );
#ifdef SO_REUSEPORT
    opt = 1;
    setsockopt( insock, SOL_SOCKET, SO_REUSEPORT, (void *) &opt,
                    sizeof( opt ) );
#endif

    /* Increase the receive buffer size to 1/2MB (8Mb/s during 1/2s) to avoid
     * packet loss caused by scheduling problems */

#ifdef SO_RCVBUF
    opt = 0x80000;
    if( setsockopt( insock, SOL_SOCKET, SO_RCVBUF, (void *) &opt,
                    sizeof( opt ) ) == -1 )
        perror( "cannot configure socket (SO_RCVBUF)");
#endif

    /* Bind it */
    memset( &inaddr , 0 , sizeof(inaddr) );
    inaddr.sin_family      = AF_INET;
    inaddr.sin_addr.s_addr = inip;
    inaddr.sin_port        = inport;

    if( bind( insock, (struct sockaddr *)&inaddr, sizeof( inaddr ) ) < 0 )
    {
        perror( "can't bind input socket" );
        return -1;
    }

    /* Join multicast group if inip is multicast address */
    if( IN_MULTICAST( ntohl(inip) ) )
    {

        if (srcip)
        {
#ifndef IP_ADD_SOURCE_MEMBERSHIP
            fprintf(stderr, "SSM not supported, trying IGMPv2\n" );
        }
#else
            struct ip_mreq_source imr;
            imr.imr_interface.s_addr = INADDR_ANY;
            imr.imr_multiaddr.s_addr = inip;
            imr.imr_sourceaddr.s_addr = srcip;

            if( setsockopt( insock, IPPROTO_IP, IP_ADD_SOURCE_MEMBERSHIP,
                        (char*)&imr, sizeof(struct ip_mreq_source) ) == -1 )
            {
                perror( "SSM join failed" );
                return -1;
            }

        } else
#endif
        {
            struct ip_mreq imr;
            imr.imr_interface.s_addr = INADDR_ANY;
            imr.imr_multiaddr.s_addr = inip;

            if( setsockopt( insock, IPPROTO_IP, IP_ADD_MEMBERSHIP,
                        (char*)&imr, sizeof(struct ip_mreq) ) == -1 )
            {
                perror( "failed to join multicast group" );
                return -1;
            }
        }
    }

    if ((outsock=socket(AF_INET,SOCK_DGRAM,0)) < 0)
    {
        perror("can't create output socket");
        return -1;
    }

    /* Increase the send buffer size to 1/2MB (8Mb/s during 1/2s) to avoid
     * packet loss caused by scheduling problems */

#ifdef SO_SNDBUF
    opt = 0x80000;
    if( setsockopt( outsock, SOL_SOCKET, SO_SNDBUF, (void *) &opt,
                    sizeof( opt ) ) == -1 )
        perror( "cannot configure socket (SO_SNDBUF)");
#endif

    /* set up destination address */
    memset( &outaddr , 0 , sizeof(outaddr) );
    outaddr.sin_family      = AF_INET;
    outaddr.sin_addr.s_addr = outip;
    outaddr.sin_port        = outport;

    /* set TTL */
    if( IN_MULTICAST( ntohl(outip) ) )
    {
        c=setsockopt(outsock,IPPROTO_IP,IP_MULTICAST_TTL,&ttl,sizeof(ttl) );
    }
    else
    {
        c=setsockopt(outsock,IPPROTO_IP,IP_TTL,&ttl,sizeof(ttl) );
    }
    if (c)
    {
        perror("can't set TTL");
        return -1;
    }

    /* Initialize RTP header */

    buf[0] = 0x80;           /* version 2 */
    buf[1] = 33;             /* MPEG-2 TS */
    opt = srcip ? srcip : inip ;
    buf[8]  = opt;      /* SSRC */
    buf[9]  = opt>>8;
    buf[10] = opt>>16;
    buf[11] = opt>>24;


    while(1)
    {
        srclen=sizeof(srcaddr);
        len=recvfrom(insock,buf+12,1500,0,(struct sockaddr*)&srcaddr, &srclen);

        if(len!=1316) continue;

        /* Ignore packets from invalid sources */
#ifndef IP_ADD_SOURCE_MEMBERSHIP
        if(srcip && srcaddr.sin_addr.s_addr != srcip) continue;   
#endif
        if(srcport && srcaddr.sin_port != srcport) continue;

        /* Update RTP sequence */
        buf[2] = rtpseq>>8;
        buf[3] = rtpseq ;
        rtpseq++;

        /* Update RTP timestamp */
        gettimeofday(&tv, NULL);
        rtpts = ( (s64)tv.tv_sec * 1000000 + (s64)tv.tv_usec ) / 10 * 9 / 10;
        buf[4] = rtpts>>24; 
        buf[5] = rtpts>>16; 
        buf[6] = rtpts>>8; 
        buf[7] = rtpts;

        sendto(outsock,buf,len+12,0,(struct sockaddr*)&outaddr,sizeof(outaddr));
    }
}