#!/usr/bin/python
# coding=utf-8
import socket
import re
import sys
import threading

UDP_IP = "0.0.0.0" # Это IP адрес машины, лучше установить 0.0.0.0
UDP_PORT = 5060 # Порт
iplist = {} # Список всех источников и количество запросов за 1 секунду в виде словаря
blacklist = [] # Черный список тех адресов, чьи запросы превысили pps
pps = 0 # Packet per second (пакеты в секунду)
blacklist_file_name = "blaclist.txt" # Имя файла BLACKLIST

def timer_callback(): # Это функция вызывается каждую секунду по тай-меру и очищает iplist
 iplist.clear() # Очистка словаря iplist

def main(): # Главная функция - обработчик
 sip_sock = socket.socket(socket.AF_INET, # Создания приемника UDP пакетов
socket.SOCK_DGRAM)

sip_sock.bind((UDP_IP, UDP_PORT)) # Прикрепление приемника UDP пакетов к 5060 порту и адресу машины

try: # Проверка на наличие входного параметра программы
 if sys.argv[1] is None: # Если его нет, то выводим ошибку и закрываем программу
  print("Please, enter value of Packet per Second as input argument")
  exit(1)
 else:
  pps = int(sys.argv[1]) # Если параметр есть - применяем его
except:
 print("Please, enter value of Packet per Second as input argument")
 exit(1)

while True: # Бесконечный цикл
 data, addr = sip_sock.recvfrom(1024) # Прием UDP пакетов

if "INVITE" or "REGISTER" in data: # Если пакет SIP - INVITE, REGISTER, то обрабатываем его
 sip_address = re.findall('From: \w+ <\w+:\w+@([\d{1-3}]+.[\d{1-3}]+.[\d{1-3}]+.[\d{1-3}]+)', data)[0] # Достаем IP адрес из поля FROM
 print(sip_address) # Вывод этого адресап

blacklist = '' # Инициализируем и очищаем blacklist
blacklist_file = open(blacklist_file_name, "wr+") # Открываем файл на считывание и запись
blacklist = str(blacklist_file.readlines()) # Считываем оттуда всю информа-цию

if sip_address in blacklist: # Если адрес в BLACKLIST, то не обрабатываем его
 print("IP {} exceeded value of packet per second, drop it".format(sip_address))
 blacklist_file.close() # Закрываем файл
continue

current_rcv_packets = iplist.get(sip_address) # Попытка получения текущего значения пришедших пакетов в секунду от данного адреса
if current_rcv_packets is not None: # Если это не первый запрос

 if current_rcv_packets >= pps: # Если количество пакектов в секунду пре-вышает pps
  print("IP address {} is exceeded value of packet per second, remove it to BLACKLIST".format(sip_address))
 blacklist.append(sip_address) # Помещаем адрес в BLACKLIST и прекра-щем работать с пакетом - дропаем
 blacklist_file.writelines(blacklist+"\n") # Записываем изменения
 blacklist_file.close() # Закрываем файл
 continue

print("{} in IP database!".format(sip_address))
current_rcv_packets = current_rcv_packets + 1 # Инече, увеличиваем счёт-чит
iplist.update(dict.fromkeys([sip_address], current_rcv_packets)) # И обновляем словарь
else:
 print("No {} IP in database, create it".format(sip_address))
 iplist.update(dict.fromkeys([sip_address], 1)) # Если это первый запррос, то добавляем адрес в словарь
 blacklist_file.close()

t = threading.Timer(1.0, timer_callback) # Инициализация таймера, который каждую секунду будет запускать функцию timer_callback

main() # Запуск основной программы
t.start() # Запуск таймера