mongoDB.py

from pymongo import MongoClient
import pyshark
import datetime
import re
import hashlib
import os

EPOCH_DATETIME = datetime.datetime(2013, 3, 28, 17, 53, 46, 769103)

class MongoHandler:

    _db = None
    _client = None
    _collection = None

    def __init__(self,collection):

        self._client = MongoClient("mongodb://localhost:27017/")
        self._db = self._client['probe_requests']
        self._collection = self._db[collection]

    def clone_db(self, old='mac_research', new="mac_research_accuracy"):
        clone_collection = self._db[old]

        if new in self._db.collection_names():
            self._db[new].drop()
        pipeline = [{"$match": {}},
                    {"$out": new},
                    ]

        clone_collection.aggregate(pipeline)



class Fingerprint(object):
    """Probe fingerprint generated by information element values."""

    def __init__(self, information_elements):
        self.values = self._extract(information_elements)

    def __repr__(self):
        """SHA256 hexdigest of specific information elements."""
        return hashlib.sha256(str(self.values).encode("UTF-8")).hexdigest()

    def _extract(self, ie):
        """Return tuple of specific information elements."""
        ordered_tag_fields = ie.get_field("wlan_mgt.tag.number").all_fields
        ordered_numbers = [int(str(field.base16_value),16) for field in ordered_tag_fields]

        return (
            ie.get_field_value("wlan_mgt.vs.ht.capabilities"), # HT capabilities info
            tuple(ordered_numbers), # Ordered list of tags numbers
            ie.get_field_value("wlan_mgt.extcap"), # Extended capabilities
            ie.get_field_value("wlan_mgt.vs.ht.ampduparam"), # HT A-MPDU parameters

            # HT MCS set bitmask
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.0to7"),
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.16to23"),
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.24to31"),
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.32"),
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.33to38"),
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.39to52"),
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.53to76"),
            ie.get_field_value("wlan_mgt.ht.mcsset.rxbitmask.8to15"),

            ie.get_field_value("wlan_mgt.supported_rates"), # Supported Rates
            ie.get_field_value("wlan_mgt.extended_supported_rates"), # Extended supported rates

            ie.get_field_value("wps.uuid_e"), # WPS UUID
            ie.get_field_value("wlan_mgt.htex.capabilities"), # HT extended capabilities
            ie.get_field_value("wlan_mgt.vs.txbf"), # HT TxBeam Forming Cap.
            ie.get_field_value("wlan_mgt.vs.asel"), # HT Antenna Selection Cap.
)


def load_pcap(pcap_location):

    pcap = pyshark.FileCapture(pcap_location,keep_packets=False,
                               display_filter="wlan.fc.type_subtype==4")
    return pcap


def add_to_db(database_info):

    db = MongoHandler(database_info['collection'])
    pcap = load_pcap(database_info['pcap_location'])
    number_of_packets = 0
    for x in pcap:
        number_of_packets +=1
        try:
            db._collection.insert_one(parse_packet_with_fingerprint(x))
        except:
            continue


def extract_ssid(val):
        """Extract SSID identifier from probe info."""

        match = re.search("(?<=SSID_)\w+", val)

        if match:
            return match.group(0)
        else:
            return None

def parse_packet_with_fingerprint(packet):
        """Parse packet into a probe dict with fingerprint."""
        dt_obj = packet.sniff_time
        mac_address = packet.wlan.sa
        return {"mac_address": mac_address,
            "timestamp": packet.sniff_timestamp,
            "ssid": extract_ssid(packet.wlan_mgt.get_field_value("wlan_mgt.ssid")),
            "fingerprint": str(Fingerprint(packet.wlan_mgt)),
            "sequence_number" : int(packet.wlan.seq)
        }



if __name__ == "__main__":

    pcap_location = '/home/alex/Documents/3rd-Year-Project/datasets/mac_info.pcap'
    collection = 'mac_info'

    database_info = {'pcap_location': pcap_location,
                     'collection' : collection}


    add_to_db(database_info)