parse3.py

import os
import shutil
import sys
import stem
import urllib.request
import geoip2.database
import tarfile

from stem.descriptor import DocumentHandler, parse_file

GEOIP_FILENAME = "./GeoLite2-City.mmdb"


def geo_ip_lookup(ip_address):
    with geoip2.database.Reader(GEOIP_FILENAME) as reader:
        try:
            response = reader.city(ip_address)
        except geoip2.errors.AddressNotFoundError:
            return False, False
        except Exception:
            print("  [-] Error while looking up %s" % ip_address)
            return False, False
        else:
            return response.location.longitude, response.location.latitude


def dl_server_descriptors(year, month):
    """ Download server descriptors from CollecTor. """
    url = "https://collector.torproject.org/archive/relay-descriptors/server-descriptors"
    filename = "server-descriptors-%s-%s.tar.xz" % (year, month)

    save_dir_path = "server-descriptors"
    if not os.path.isdir(save_dir_path):
        os.mkdir(save_dir_path)
    save_path = "%s/%s" % (save_dir_path, filename)
    # Check if the archive exists.
    if os.path.isfile(save_path):
        print("  [+] Server descriptors %s found" % save_path)
        return save_path
    # Check if the directory (uncompressed) exists.
    if os.path.isdir("%s" % (save_path[:-7])):
        print("  [+] Server descriptors %s found" % (save_path[:-7]))
        return save_path

    print("  [+] Downloading server descriptors %s/%s" % (url, filename))
    try:
        request = urllib.request.urlopen("%s/%s" % (url, filename))
        if request.code != 200:
            print("  [-] Unable to fetch server descriptors %s at %s" %
                  (filename, url))
            return None
    except Exception:
        print("  [-] Unable to fetch %s/%s" % (url, filename))
        return None
    fp = open(save_path, "wb+")
    fp.write(request.read())
    fp.close()
    return save_path


def dl_consensus(year, month):
    """ Download consensus from CollecTor. """
    url = "https://collector.torproject.org/archive/relay-descriptors/consensuses"
    filename = "consensuses-%s-%s.tar.xz" % (year, month)

    save_dir_path = "consensuses"
    if not os.path.isdir(save_dir_path):
        os.mkdir(save_dir_path)
    save_path = "%s/%s" % (save_dir_path, filename)
    # Check if the archive exists.
    if os.path.isfile(save_path):
        print("  [+] Consensus %s found" % save_path)
        return save_path
    # Check if the directory (uncompressed) exists.
    if os.path.isdir("%s" % (save_path[:-7])):
        print("  [+] Consensus %s found" % (save_path[:-7]))
        return save_path

    print("  [+] Downloading consensus %s/%s" % (url, filename))
    try:
        request = urllib.request.urlopen("%s/%s" % (url, filename))
        if request.code != 200:
            print("  [-] Unable to fetch consensus %s at %s" % (filename, url))
            return None
    except Exception:
        print("  [-] Unable to fetch %s/%s" % (url, filename))
        return None

    fp = open(save_path, "wb+")
    fp.write(request.read())
    fp.close()
    return save_path


def dl_extra_infos(year, month):
    """ Download extra infos from CollecTor. """
    url = "https://collector.torproject.org/archive/relay-descriptors/extra-infos"
    filename = "extra-infos-%s-%s.tar.xz" % (year, month)

    save_dir_path = "extra-infos"
    if not os.path.isdir(save_dir_path):
        os.mkdir(save_dir_path)
    save_path = "%s/%s" % (save_dir_path, filename)
    # Check if the archive exists.
    if os.path.isfile(save_path):
        print("  [+] Extra infos %s found" % save_path)
        return save_path
    # Check if the directory (uncompressed) exists.
    if os.path.isdir("%s" % (save_path[:-7])):
        print("  [+] Extra infos %s found" % (save_path[:-7]))
        return save_path

    print("  [+] Downloading extra infos %s/%s" % (url, filename))
    try:
        request = urllib.request.urlopen("%s/%s" % (url, filename))
        if request.code != 200:
            print("  [-] Unable to fetch extra infos %s at %s" %
                  (filename, url))
            return None
    except Exception:
        print("  [-] Unable to fetch %s/%s" % (url, filename))
        return None
    fp = open(save_path, "wb+")
    fp.write(request.read())
    fp.close()
    return save_path


def uncompress(path, dst):
    # Remove .tar.xz
    dirname = path[:-7]
    if os.path.isdir(dirname):
        print("  [+] %s Already uncompressed into %s/%s" % (path, dst, dirname))
        return
    print("  [+] Uncompressing %s into %s/%s" % (path, dst, dirname))
    with tarfile.open(path) as f:
        f.extractall(dst)


def get_previous_data(year, month, day):
    # If day is undefined or if day is 1, we have to get the previous month
    # server descriptors data to get the descriptors.
    prev_sd_path = prev_ei_path = None
    if day == 0 or day == 1:
        prev_year = year
        prev_month = month
        if month == 1:
            prev_year -= 1
            prev_month = 12
        else:
            prev_month -= 1
        str_month = str(prev_month)
        if prev_month < 10:
            str_month = "0%d" % prev_month
        prev_sd_path = dl_server_descriptors(prev_year, str_month)
        prev_ei_path = dl_extra_infos(prev_year, str_month)
    return prev_sd_path, prev_ei_path


def create_csv_file(year, month, day):
    # Process the consensuses that we are interested in.
    csv_filename = 'data/relays-%s-%s-%s-00-00-00.csv' % \
        (year, month, day)
    if os.path.exists(csv_filename):
        print("  [+] CSV %s exists, skipping!" % csv_filename)
        return None
    csv = open(csv_filename, 'w+')
    print("  [+] Creating CSV file %s" % csv_filename)
    csv.write(
        'Name,Fingerprint,Flags,IP,OrPort,ObservedBW,GuardClients,DirClients,Uptime,Longitude,Latitude)\n')
    return csv


def client_ips_to_string(ei_dict, sep):
    List = []
    for key, value in ei_dict.items():
        List.append('%s:%s' % (key, value))
    return sep.join(List)


def write_csv_data(consensus, sd_path, prev_sd_path, ei_path, prev_ei_path, year, month, day):
    """ Write data from consensus to CSV file """
    csv_fp = create_csv_file(year, month, day)
    if csv_fp is None:
        # CSV file already exists.
        return None

    for desc in consensus.routers.values():
        # Check for longitude and latitude. Without this, the entry is useless.
        lon, lat = geo_ip_lookup(desc.address)
        if lon is False and lat is False or lon is None or lat is None:
            continue

        fp = desc.fingerprint
        digest = desc.digest.lower()
        sd_filename = "%s/%s/%s/%s" % (sd_path[:-7], digest[0], digest[1], digest)
        try:
            sd = next(parse_file(sd_filename))
        except Exception:
            if prev_sd_path is None:
                continue
            sd_filename = "%s/%s/%s/%s" % (
                prev_sd_path[:-7], digest[0], digest[1], digest)
            try:
                sd = next(parse_file(sd_filename))
            except Exception:
                print("  [-] Server descriptor %s not found" % digest)
                continue

        # Open extra info.
        entry_ips = ""
        dir_ips = ""
        if sd.extra_info_digest is not None:
            digest = sd.extra_info_digest.lower()
            ei_filename = "%s/%s/%s/%s" % (ei_path[:-7],
                                           digest[0], digest[1], digest)
            try:
                ei = next(parse_file(ei_filename))
            except Exception:
                if prev_ei_path is None:
                    continue
                ei_filename = "%s/%s/%s/%s" % (
                    prev_ei_path[:-7], digest[0], digest[1], digest)
                try:
                    ei = next(parse_file(ei_filename))
                except Exception:
                    print("  [-] Extra info %s not found" % ei_filename)
                    continue
            try:
                # Any Guard client ips?
                if ei.entry_ips is not None and len(ei.entry_ips) != 0:
                    entry_ips = client_ips_to_string(ei.entry_ips, "|")
            except Exception:
                pass
            try:
                # Any Directory client ips?
                if ei.dir_v3_requests is not None and len(ei.dir_v3_requests) != 0:
                    dir_ips = client_ips_to_string(ei.dir_v3_requests, "|")
            except Exception:
                pass

        # Get relay flags.
        flag = "M"
        if stem.Flag.GUARD in desc.flags:
            flag += "G"
        if stem.Flag.EXIT in desc.flags:
            flag += "E"
        if stem.Flag.HSDIR in desc.flags:
            flag += "H"

        csv_fp.write("%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s\n" % (desc.nickname,
                                                             desc.fingerprint, flag, desc.address, desc.or_port,
                                                             float(
                                                                 sd.observed_bandwidth/1000.0/1000.0), entry_ips,
                                                             dir_ips, sd.uptime, lon, lat))
    csv_fp.close()


def make_monthly_csv(year, month, day):
    """
    Create the CSV files for the given year/month. If day is defined, only
    create the file for that day else all the day at midnight.
    """
    match_found = False
    str_month = str(month)
    if month < 10:
        str_month = "0%d" % month
    consensus_path = dl_consensus(year, str_month)
    if consensus_path is None:
        return None
    sd_path = dl_server_descriptors(year, str_month)
    ei_path = dl_extra_infos(year, str_month)
    if sd_path is None or ei_path is None:
        print("Unable to create CSV files for %s-%s" % (year, str_month))
        return None
    prev_sd_path, prev_ei_path = get_previous_data(year, month, day)
    if prev_sd_path is not None:
        uncompress(prev_sd_path, './server-descriptors')
        uncompress(prev_ei_path, './extra-infos')
    uncompress(consensus_path, './consensuses')
    uncompress(sd_path, './server-descriptors')
    uncompress(ei_path, './extra-infos')
    # We have the data, let's create the csv files for the requested date.
    for dir_day in os.listdir('./%s' % (consensus_path[:-7])):
        str_day = str(day)
        if day < 10:
            str_day = "0%d" % day
        if day != 0 and str_day != dir_day:
            continue
        match_found = True
        consensus_pathname = \
            "./consensuses/consensuses-%s-%s/%s/%s-%s-%s-00-00-00-consensus" % \
            (year, str_month, dir_day, year, str_month, dir_day)
        print("  [+] Reading consensus %s" % consensus_pathname)
        try:
            consensus = next(parse_file(consensus_pathname,
                             document_handler=DocumentHandler.DOCUMENT))
        except Exception:
            print("  [-] Consensus %s not found. Skipping!" % consensus_pathname)
            continue

        # Nullify the previous path if we aren't the first of the month.
        if dir_day != "01":
            prev_ei_path = None
            prev_sd_path = None
        write_csv_data(consensus, sd_path, prev_sd_path, ei_path, prev_ei_path,
                       str(year), str_month, dir_day)

    if match_found is False:
        print("  [-] Date not found in consensus")
        shutil.rmtree(consensus_path)
        shutil.rmtree(sd_path)
        if prev_sd_path is not None:
            shutil.rmtree(prev_sd_path)


def make_yearly_csv(year):
    """ DOC TBC """
    for month in range(1, 13):
        make_monthly_csv(year, month, 0)


def run(year, month, day):
    """
    Using the given date, download the needed files and create the csv file(s).
    """
    if month != 0:
        make_monthly_csv(year, month, day)
    else:
        make_yearly_csv(year)

    # Cleanup what's left if any.
    # for dirname in os.listdir('./consensuses'):
    #    shutil.rmtree(dirname)
    # for dirname in os.listdir('./server-descriptors'):
    #    shutil.rmtree(dirname)


def usage():
    print("Usage: %s <YEAR> [<MONTH> [<DAY>]]" % (sys.argv[0]))
    sys.exit(1)


if __name__ == '__main__':
    if len(sys.argv) == 1:
        usage()

    # Make sure we have a GeoIP database (maxmind)
    if not os.path.isfile(GEOIP_FILENAME):
        print("%s not found. It must be in the same directory as this script." %
              GEOIP_FILENAME)
        print("Get the Maxmind GeoIP2 city database here:")
        print("-> https://dev.maxmind.com/geoip/geolite2-free-geolocation-data")
        sys.exit(1)

    month = day = 0
    try:
        year = int(sys.argv[1])
        if len(sys.argv) > 2:
            month = int(sys.argv[2])
        if len(sys.argv) > 3:
            day = int(sys.argv[3])
        # Create the data repository
        if not os.path.isdir("./data"):
            os.mkdir("./data")
        run(year, month, day)
    except ValueError as e:
        print("Invalid argument.")
        usage()