#! /usr/bin/env python
#coding=utf-8
\'\'\'
fileName: server.py
新增功能：决定客户端发送周期
\'\'\'
from tornado.tcpserver import TCPServer  
from tornado.ioloop  import IOLoop 
import pymongo 
import json

from Queue import Queue
import threading
import time

# 处理队列
processdata=Queue()

# ChatServer类与客户端通信，并获取客户端发来的data
class ChatServer(TCPServer):  
    def handle_stream(self, stream, address): 
        print "New connection :", address, stream 
        Connection(stream, address) 
        print "connection num is:", len(Connection.clients)

class Connection(object):  
    clients = set()  
    def __init__(self, stream, address): 
        Connection.clients.add(self) 
        self._stream = stream  
        self._address = address  
        self._stream.set_close_callback(self.on_close)  
        self.read_message()  
    
    # 读取客户端发来的信息  
    def read_message(self):  
        self._stream.read_until(\'\n\', self.broadcast_messages)  

    # 处理信息
    def broadcast_messages(self, data):  
        # print \'client said\',data[:-1]
        if data[:-1] == \'cycle\':    # 客户端请求周期
            cycle = \'60\'            # 自定义客户端发送周期
            for conn in Connection.clients:  
                conn.send_message(cycle)
            self.read_message()
        else:       # 正常接收监控数据
            for conn in Connection.clients:  
                conn.send_message(\'Server have received the data successfully!\')
            self.read_message()  
            #存储到mongodb数据库
            if data[0] == \'{\':
                self.handleData(data)

    def handleData(self,data):
        conn = pymongo.Connection("localhost", 27017)
        db = conn[\'networkSecurity\']
        systemInfo = db[\'systemInfo\']
        info = json.loads(data[:-1])
        # 构造dict数据
        message = {\'IP\':info[u\'IP\'],\'CPUstate\':info[u\'CPUstate\'],\'Memorystate\':info[u\'Memorystate\'],
            \'PortState\':info[u\'PortState\'],\'ProcessName\':info[u\'ProcessName\']}
        print \'Client said :\nIP:%s\nCPUstate:%s\nMemorystate:%s\nPortState:%s\nProcessName:%s\'%(message[\'IP\'],message[\'CPUstate\'],message[\'Memorystate\'],message[\'PortState\'],message[\'ProcessName\'])
        # 将数据放入队列
        processdata.put(message)
        # 将数据存入数据库
        systemInfo.insert(message)
        print \'success to store the data!\'

    def send_message(self, data):  
        self._stream.write(data) 
        # self._stream.write(\'yessssssssssssss\') 
          
    def on_close(self):  
        print "A user has left the chat room.", self._address
        Connection.clients.remove(self)  

#消费者
class Consumer(threading.Thread):
    def __init__(self, t_name):
        threading.Thread.__init__(self, name=t_name)
    def run(self):
        print "%s: %s is consuming in the queue!/n" %(time.ctime(), self.getName())
        message = processdata.get()
        # print \'the message in the queue is :   \',message
        # print type(message)
        monitorSystem(message)

# 黑白名单匹配，info为字典
def monitorSystem(info):
    warning = 0
    whiteList = [\'cmd.exe\']
    blackList = [\'sublime_text.exe\']
    # for info in systemInfo.find():
    #     print info
    IP = info[\'IP\']
    processName = info[\'ProcessName\']
    for process in processName:
        if process in blackList:
            warning = 1
            print \'Process  %s in black list is running in IP %s ! \'%(process,IP)
    for process in whiteList:
        if process not in processName:
            warning = 1
            print \'Process %s in white list is not running in IP %s ! \'%(process,IP)
    if warning == 0:
        print \'Host %s is running legally ! \'%IP

def main():
    # 启动服务器
    # 生产进程：接受数据
    print "Server start ......"  
    server = ChatServer()  
    server.listen(8000)  
    # 消费进程：处理数据，黑白名单匹配
    consumer = Consumer(\'Con.\')
    consumer.start()

    IOLoop.instance().start()
 
if __name__ == \'__main__\':  
    main()
 

#! /usr/bin/env python
#coding=utf-8
\'\'\'
fileName：client.py
监控windows信息：CPU占有率，内存占有率，端口开放情况，当前进程名称
数据格式：    {\'IP\':getIp(),\'CPUstate\':getCPUState(),\'Memorystate\':getMemoryState(),
            \'PortState\':getPortState(),\'ProcessName\':getProcessName()}
新增功能：    接收服务器定义的数据发送周期
\'\'\'
import socket
import time
import psutil
import json
import thread


# 启动服务器 
def runServer(s,HOST,PORT):
    s.connect((HOST, PORT))
    print \'connect success!\'

# 关闭服务器
def closeServer(s):
    time.sleep(1)
    s.close()

# 向服务器发送数据
def sendMessage(s,portState):
    s.send(json.dumps(packMessage(portState)))
    time.sleep(1)
    s.sendall(\'\n\')

# 从服务器接收数据
def recvMessage(s):
    data = s.recv(1024)
    print \'Received\', repr(data)

# 构造数据包
def packMessage(portState):
    message =     {\'IP\':getIp(),\'CPUstate\':getCPUState(),\'Memorystate\':getMemoryState(),
                \'PortState\':portState,\'ProcessName\':getProcessName()}
    print \'My message is :\nIP:%s\nCPUstate:%s\nMemorystate:%s\nPortState:%s\nProcessName:%s\'%(message[\'IP\'],message[\'CPUstate\'],message[\'Memorystate\'],message[\'PortState\'],message[\'ProcessName\'])
    return message

# 获取本机IP
def getIp():
    myname = socket.getfqdn(socket.gethostname())
    myaddr = socket.gethostbyname(myname)
    return myaddr

# 获取CPU使用率 
def getCPUState(interval=1):
    return (str(psutil.cpu_percent(interval)) + "%")

# 获取内存使用率  
def getMemoryState():
    mem_rate = 0
    for pnum in psutil.pids():
        p = psutil.Process(pnum) 
        mem_rate = mem_rate + p.memory_percent()
    return "%.2f%%"%mem_rate

# 输入IP和端口号，扫描判断端口是否开放
def socket_port(ip,port,portList):
    try:
        if port >= 65535:
            print u\'端口扫描结束\'
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        result = s.connect_ex((ip,port))
        if result == 0:
            lock.acquire()
            portList.append(port)
            lock.release()
        s.close()
    except:
        print u\'端口扫描异常\'

# 输入IP，扫描IP的0-65534端口情况 
def ip_scan(ip):
    portList = []
    socket.setdefaulttimeout(3)
    try:
        for i in range(0,65534):
            thread.start_new_thread(socket_port,(ip,int(i),portList))
        # 返回所有开放的端口号
        return portList
    except:
        print u\'扫描ip出错\'

# 获取正在运行的进程名称
def getProcessName():
    ProcessNameList = []
    for pnum in psutil.pids():
        p = psutil.Process(pnum) 
        ProcessNameList.append(p.name())
    return ProcessNameList

if __name__ == "__main__" :
    HOST = \'192.168.111.130\'    # 服务器IP
    PORT = 8000                    # 端口号
    myIP = getIp()

    # 接收服务器定义的数据发送周期
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((HOST, PORT))
    s.sendall(\'cycle\n\')
    data = s.recv(1024)
    print \'my cycle is\', data
    cycle=int(data)

    while True:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        runServer(s,HOST,PORT)
        lock = thread.allocate_lock()
        portState = ip_scan(myIP)
        sendMessage(s,portState)
        closeServer(s)
        time.sleep(cycle)