Onionr/static-data/default-plugins/torgossip/peerdb.py

"""Onionr - Private P2P Communication.

Torgossip peer safedb interface
"""
from base64 import b32decode
from struct import unpack, pack
from time import time
from typing import List
import binascii

from utils.identifyhome import identify_home
import safedb
"""
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""


class TorGossipPeers:  # name it this way to avoid collisions in SharedState

    PACK_FORMAT = "qQ"

    def __init__(self, encrypt=False):
        self.db = safedb.SafeDB(
            identify_home() + "/torgossip-peers.db", protected=encrypt)

    def _shrink_peer_address(self, peer):
        # strip .onion and b32decode peer address for lower database mem usage
        if len(peer) == 62:  # constant time
            return b32decode(peer[:-6])  # O(N)
        else:
            try:
                return b32decode(peer.upper())
            except binascii.Error:
                pass
        return peer

    def _pack_and_store_info(self, peer, new_score=None, new_seen=None):
        # Repack peer information with new value(s) and store it
        peer = self._shrink_peer_address(peer)
        try:
            score, seen = unpack(self.PACK_FORMAT, self.db.get(peer))
        except KeyError:
            score = seen = 0

        if new_score:
            score = new_score
        if new_seen:
            seen = new_seen

        self.db.put(peer, pack(self.PACK_FORMAT, score, seen))

    def get_highest_score_peers(self, max) -> List:
        assert max >= 1
        peer = self.db.db_conn.firstkey()
        if peer == b'enc':
            peer = self.db.db_conn.nextkey(peer)

        if not peer:
            return []

        assert len(peer) == 35

        top = [(peer, self.db.get(peer))]

        while peer:
            peer = self.db.db_conn.nextkey(peer)
            if not peer:
                break
            if peer == b"enc":
                continue
            # score, seen is data
            peer_data = self.db.get(peer)
            overwrite = None

            if len(top) != max:
                top.append((peer, peer_data))
                peer = self.db.db_conn.nextkey(peer)
                continue
            for count, cur_top in enumerate(top):
                # if peer score is greater than any set peer, overwrite
                if unpack(self.PACK_FORMAT, cur_top[1])[0] < \
                        unpack(self.PACK_FORMAT, peer_data)[0]:
                    overwrite = count
                    break  # below else won't execute, so it will be overwritten
            else:
                # if not overwriting, go to next peer
                peer = self.db.db_conn.nextkey(peer)
                continue
            top[overwrite] = (peer, peer_data)
        return top

    def add_score(self, peer, plus):
        shrunk = self._shrink_peer_address(peer)
        score, _ = unpack("qQ", self.db.get(shrunk))

        self._pack_and_store_info(
            shrunk, new_score=score + plus)

    def update_seen(self, peer):
        peer = self._shrink_peer_address(peer)
        self._pack_and_store_info(peer, new_seen=int(time()))

    def add_peer(self, peer):
        self.update_seen(peer)

    def remove_peer(self, peer):
        peer = self._shrink_peer_address(peer)
        del self.db.db_conn[peer]