Saves inv messages from pcap files in Redis

Author: ayeowch

Diff for: pcap.conf

@@ -0,0 +1,10 @@
+[pcap]
+
+# Logfile
+logfile = pcap.log
+
+# Print debug output
+debug = False
+
+# Relative path to directory containing pcap files
+pcap_dir = data/pcap

Diff for: pcap.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# pcap.py - Saves inv messages from pcap files in Redis.
+#
+# Copyright (c) 2014 Addy Yeow Chin Heng <[email protected]>
+#
+# Permission is hereby granted, free of charge, to any person obtaining
+# a copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and to
+# permit persons to whom the Software is furnished to do so, subject to
+# the following conditions:
+#
+# The above copyright notice and this permission notice shall be
+# included in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+"""
+Saves inv messages from pcap files in Redis.
+"""
+
+import dpkt
+import glob
+import logging
+import os
+import redis
+import socket
+import sys
+import threading
+import time
+from ConfigParser import ConfigParser
+
+from protocol import ProtocolError, Serializer
+
+# Redis connection setup
+REDIS_SOCKET = os.environ.get('REDIS_SOCKET', "/tmp/redis.sock")
+REDIS_PASSWORD = os.environ.get('REDIS_PASSWORD', None)
+REDIS_CONN = redis.StrictRedis(unix_socket_path=REDIS_SOCKET,
+                               password=REDIS_PASSWORD)
+
+SETTINGS = {}
+
+
+def save_invs(timestamp, node, invs):
+    """
+    Adds inv messages into the inv set in Redis.
+    """
+    timestamp = int(timestamp * 1000)  # in ms
+    redis_pipe = REDIS_CONN.pipeline()
+    for inv in invs:
+        logging.debug("[{}] {}:{}".format(timestamp, inv['type'], inv['hash']))
+        key = "inv:{}:{}".format(inv['type'], inv['hash'])
+        redis_pipe.zadd(key, timestamp, node)
+        redis_pipe.expire(key, 28800)  # Expires in 8 hours
+    redis_pipe.execute()
+
+
+def get_invs(filepath):
+    """
+    Extracts inv messages from the specified pcap file.
+    """
+    serializer = Serializer()
+    pcap_file = open(filepath)
+    pcap_reader = dpkt.pcap.Reader(pcap_file)
+    for timestamp, buf in pcap_reader:
+        frame = dpkt.ethernet.Ethernet(buf)
+        ip_packet = frame.data
+        if isinstance(ip_packet.data, dpkt.tcp.TCP):
+            tcp_packet = ip_packet.data
+            payload = tcp_packet.data
+            if len(payload) > 0:
+                try:
+                    (msg, _) = serializer.deserialize_msg(payload)
+                except ProtocolError as err:
+                    pass
+                else:
+                    if msg['command'] == "inv":
+                        if ip_packet.v == 6:
+                            address = socket.inet_ntop(socket.AF_INET6,
+                                                       ip_packet.src)
+                        else:
+                            address = socket.inet_ntop(socket.AF_INET,
+                                                       ip_packet.src)
+                        node = (address, tcp_packet.sport)
+                        save_invs(timestamp, node, msg['inventory'])
+    pcap_file.close()
+
+
+def cron():
+    """
+    Periodically fetches oldest pcap file to extract inv messages from.
+    """
+    while True:
+        time.sleep(5)
+
+        try:
+            oldest = min(glob.iglob("{}/*.pcap".format(SETTINGS['pcap_dir'])))
+        except ValueError as err:
+            logging.warning(err)
+            continue
+        latest = max(glob.iglob("{}/*.pcap".format(SETTINGS['pcap_dir'])))
+        if oldest == latest:
+            continue
+        dump = oldest
+        logging.info("Dump: {}".format(dump))
+
+        start = time.time()
+        get_invs(dump)
+        end = time.time()
+        elapsed = end - start
+        logging.info("Elapsed: {}".format(elapsed))
+
+        os.remove(dump)
+
+
+def init_settings(argv):
+    """
+    Populates SETTINGS with key-value pairs from configuration file.
+    """
+    conf = ConfigParser()
+    conf.read(argv[1])
+    SETTINGS['logfile'] = conf.get('pcap', 'logfile')
+    SETTINGS['debug'] = conf.getboolean('pcap', 'debug')
+    SETTINGS['pcap_dir'] = conf.get('pcap', 'pcap_dir')
+    if not os.path.exists(SETTINGS['pcap_dir']):
+        os.makedirs(SETTINGS['pcap_dir'])
+
+
+def main(argv):
+    if len(argv) < 2 or not os.path.exists(argv[1]):
+        print("Usage: pcap.py [config]")
+        return 1
+
+    # Initialize global settings
+    init_settings(argv)
+
+    # Initialize logger
+    loglevel = logging.INFO
+    if SETTINGS['debug']:
+        loglevel = logging.DEBUG
+
+    logformat = ("%(asctime)s,%(msecs)05.1f %(levelname)s (%(funcName)s) "
+                 "%(message)s")
+    logging.basicConfig(level=loglevel,
+                        format=logformat,
+                        filename=SETTINGS['logfile'],
+                        filemode='w')
+    print("Writing output to {}, press CTRL+C to terminate..".format(
+          SETTINGS['logfile']))
+
+    threading.Thread(target=cron).start()
+
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv))

Diff for: protocol.py

@@ -272,8 +272,12 @@ def serialize_ping_payload(self, nonce):
 
     def deserialize_ping_payload(self, data):
         data = StringIO(data)
+        try:
+            nonce = struct.unpack("<Q", data.read(8))[0]
+        except struct.error as err:
+            raise ReadError(err)
         msg = {
-            'nonce': struct.unpack("<Q", data.read(8))[0],
+            'nonce': nonce,
         }
         return msg
 

Diff for: requirements.txt

@@ -1,2 +1,3 @@
+-e svn+http://dpkt.googlecode.com/svn/trunk@90#egg=dpkt-1.8-py2.7-dev_r90
 gevent==1.0.1
 redis==2.9.1

Diff for: start.sh

@@ -16,3 +16,5 @@ python -u export.py export.conf > export.out 2>&1 &
 python -u chart.py chart.conf > chart.out 2>&1 &
 
 python -u seeder.py seeder.conf > seeder.out 2>&1 &
+
+python -u pcap.py pcap.conf > pcap.out 2>&1 &