udp_ring_dump.py

#!/usr/bin/env python
#coding=utf-8

import os
import sys
import getopt
import re
import subprocess
import time
import socket
import struct
import select
import traceback
import glob

version = "0.1.2";
chname = "";
churi = "";
chroot = "";
MAX_SIZE = 0x200000;
MAX_COUNT = 20;
files_count = 0;
has_save_status = False;
log_file_path = "";

#226.6.7.1:5001
def usage():
    info = """Usage Info:
          -h | --help:    Get help infomation.
          -v | --version: Get tool's current version.
          --uri:    Set the source address.
          --name:   Set the channel's name.
          --root:   Set dump-files' path
          --msize:  (MB)Limited the single-file's MAX size[default:%d MB].
          --mcount: Limited the files' MAX count[default:%d].
          
          eg.
              src_dump --name=xxx --uri=udp://xxx.xxx.xxx.xxx:xxxx --root=/xxx
          """ %(MAX_SIZE/0x100000, MAX_COUNT);
    print info;


def uri_parse(uri):
    uri = uri.strip();
    protocol = "";
    ip = "";
    port = 0;
    
    m = re.match(r"(\w+?)://(\d+?\.\d+?\.\d+?\.\d+?):(\d+)", uri);
    if m:
        if len(m.groups()) == 3:
            protocol = m.group(1);
            ip = m.group(2);
            port = int(m.group(3));
        else:
            print "URI format Error!!!"
    else:
        print "URI format Error!!!"

    return protocol, ip, port;

def check_paramters(protocol, ip, port):
    if protocol == "":
        print "** ERROR: Invalid uri !";
        return False;
    elif protocol != "udp":
        print "** ERROR: unsupport protocol !"
        return False;
    if ip == "":
        print "** ERROR: Invalid IP !";
        return False;
    else:
        ips = ip.split(".");
        ips = [int(x) for x in ips];
        for x in ips:
            if (x < 0 or x > 255):
                print "** ERROR: Invalid IP !";
                return False;
    if port < 0 and port > 65535:
        print "** ERROR: Invalid port !";
        return False;

    if chname == "":
        print "** ERROR: Invalid name !"
        return False;

    if chroot == "":
        print "** ERROR: Invalid root !"
        return False;
    return True;

def get_time_string():
    return time.strftime('[%Y%m%d-%02H:%02M:%02S]').strip();

def make_filename():
    date = time.strftime('%Y%m%d%02H%02M%02S').strip();
    filename = "%s/%s/%s_%s.ts" %(chroot, chname, date, chname);
    filename = filename.replace("//", "/");
    return filename;

def try_clean_dir(dir_path):
    cmd = """ls %s/*.ts""" % (dir_path);
    pipe = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True);
    lines = pipe.stdout.readlines();
    pipe.wait();
    if(pipe.stdout):
        pipe.stdout.close();

    lines = [x.strip() for x in lines];
    if(len(lines) > MAX_COUNT):
        del_name = ""
        ctime_min = 0;
        for line in lines:
            statinfo=os.stat(line)
            if (ctime_min == 0 or ctime_min > statinfo.st_ctime):
                ctime_min = statinfo.st_ctime;
                del_name = line;
            
        global files_count;
        files_count -= 1;
        cmd = "rm -f %s" %(del_name);
        
        if (os.system(cmd) != 0):
            print get_time_string(), cmd, "FAILED !"
        else:
            print get_time_string(), "[%d]%s" %(files_count, cmd);

def ip2int(ip):
    ips = ip.split(".");
    ips = [int(x) for x in ips];

    return ((ips[0] << 24) | (ips[1] << 16) | (ips[2] << 8) | (ips[3]) );

def save_status(dir_path):
    fp = open(log_file_path, "w+");

    padding = "    ";
    
    date = get_time_string();
    fp.write(date + "Recv TIMEOUT !\n");
    
    cmd = """ls %s/*.ts""" % (dir_path);
    pipe = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True);
    lines = pipe.stdout.readlines();
    pipe.wait();
    if(pipe.stdout):
        pipe.stdout.close();
    lines = [x.strip() for x in lines];

    autosave_dir = "%s/autosave" %(dir_path);
    cmd = "mkdir -p " + autosave_dir;
    if(os.system(cmd) != 0):
        info = padding + cmd + " FAILED!!!";
        fp.write(info + "\n");
        print info;
    
    for line in lines:
        cmd = "mv %s %s" %(line, autosave_dir);
        info = padding + cmd;
        fp.write(info + "\n");
        print info;
        if(os.system(cmd) != 0):
            info = padding + cmd + " FAILED!!!";
            fp.write(info + "\n");
            print info;
    
    fp.close();

def get_dir_files_count(dir_path, file_type):
    pathname = dir_path + "/" + file_type;
    lines = glob.glob(pathname);
    return len(lines);
        

if __name__ == "__main__":
    if (len(sys.argv) < 2):
        usage();
        exit(0);
        
    opts, args = getopt.getopt(sys.argv[1:],'vh', ['help', 'version', 'uri=', 'name=', "root=", "msize=", "mcount="]);
    for key,value in opts:
        if key in ("-v", "--version"):
            print "Current Version:", version;
            exit(0);
        elif key in ("-h", "--help" ):
            usage();
            exit(0);
        elif key in ("--uri"):
            churi = value;
        elif key in ("--name"):
            chname = value;
        elif key in ("--root"):
            chroot  = value;
        elif key in ("--msize"):
            MAX_SIZE  = int(value) * 0x100000;
        elif key in ("--mcount"):
            MAX_COUNT = int(value);
        else:
            print "UNKONOWN args:" + str(key);


    protocol, ip, port = uri_parse(churi);

    if not check_paramters(protocol, ip, port):
        usage();
        exit(1);  
    

    cmd = """mkdir -p %s/%s
          """ %(chroot, chname);
    if (os.system(cmd) != 0):
        print "Create Dir FAILED !";
        exit(0);

    s = socket.socket(socket.AF_INET , socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1);

    try:
        if (ip2int(ip) >= ip2int("224.0.0.0")) and (ip2int(ip) <= ip2int("239.255.255.255")):
            s.bind(('' , port));  
            mreq = struct.pack("=4sl" , socket.inet_aton(ip) , socket.INADDR_ANY)  
            s.setsockopt(socket.IPPROTO_IP , socket.IP_ADD_MEMBERSHIP , mreq)
        else:
            s.bind((ip, port));
    except socket.error:
        print "Socket bind failed !"
        traceback.print_exc();
        exit(0)

##    socket.setblocking(1);
    s.settimeout(5) 
    
    while True:
        has_save_satus = False;
        nb_recv = 0;
        filename = make_filename();
        fp = open(filename, "wb");        
        
        dir_path = "%s/%s" %(chroot, chname);
        dir_path = dir_path.replace("//", "/");
        log_file_path = "%s/%s.log" %(dir_path, chname);
        log_file_path = log_file_path.replace("//", "/");

        

        files_count = get_dir_files_count(dir_path, "*.ts");
        print get_time_string(), "[%d]Writing:%s" % (files_count, filename);
        
        try_clean_dir(dir_path);

        while nb_recv < MAX_SIZE:
            try:
                data = s.recv(4096)
            except socket.timeout: #FIXME
                info = get_time_string() + "Time Out try it again."
                
                if not has_save_status:
                    has_save_status = True;
                    save_status(dir_path);

                cmd = '''echo "%s" >> %s ''' % (info, log_file_path);
                print info;
                os.system(cmd);
                continue;
            nb_recv += len(data)
            fp.write(data);
        fp.close();

    s.close()