async def verify_seed(self, message: BARMessage, bn: BootstrapIdentity):
n_of_peers = PeerView.get_total_peers_per_epoch(message.token.epoch, bn.id)
while n_of_peers == 0:
Logger.get_instance().debug_item('Waiting for view message ...')
await asyncio.sleep(0.5)
n_of_peers = PeerView.get_total_peers_per_epoch(message.token.epoch, bn.id)
partners_index = list(reversed(self.crypto.get_random().prng(message.token.bn_signature, n_of_peers - 1,
MAX_CONTACTING_PEERS * self.RETRY)))
my_pk = self.crypto.get_ec().dump_public_key(self.crypto.get_ec().public_key)
if not my_pk == self.crypto.get_ec().dump_public_key(message.to_peer.public_key):
return False
# if i am the one requesting the exchange
current_epoch_tokens = Token.find_all_tokens(message.token.epoch)
if any(token == message.token.bn_signature for token in current_epoch_tokens):
return True
while len(partners_index) > 0:
p_index = partners_index.pop()
partner = PeerView.get_partner(p_index)
if partner.public_key == self.crypto.get_ec().dump_public_key(message.from_peer.public_key):
continue
elif partner.public_key == self.crypto.get_ec().dump_public_key(
message.to_peer.public_key) and partner.is_me:
return True
elif partner.public_key != self.crypto.get_ec().dump_public_key(
message.to_peer.public_key) and not partner.is_me:
return False
async def run(self, reader: StreamReader,
writer: StreamWriter) -> NoReturn:
Logger.get_instance().debug_item('New connection from {}'.format(
writer.get_extra_info('peername')))
self.add_connection(writer)
while True:
try:
msgs_bytes = await self.wait_message(reader)
reader._eof = False
# if len(msgs_bytes) == 1 and msgs_bytes[0] is self.EMPTY_BYTES:
# reader._eof = True
# await self.drop_connection(writer)
# Logger.get_instance().debug_item(
# 'Connection with {} has been dropped'.format(writer.get_extra_info('peername')))
# break
for msg in msgs_bytes:
await self.handle(msg, writer)
except Exception as e:
traceback.print_exc()
Logger.get_instance().debug_item(
'An error has occurred: {}'.format(e.args[0]),
LogLevels.ERROR)
await PubSub().remove_all()
await self.drop_connection(writer)
except KeyboardInterrupt:
return
async def _handle(self, connection: StreamWriter,
message: PromiseBARMessage):
if not await self.is_valid_message(message):
Logger.get_instance().debug_item(
'Invalid request... sending PoM 1')
await self.send_pom(Misbehaviour.BAD_SEED, message, connection)
return
ser_needed, ser_promised = json.dumps(message.needed), json.dumps(
message.promised)
valid_promise = self.is_valid_promise(ser_needed, ser_promised,
message.token.bn_signature)
if not valid_promise:
Logger.get_instance().debug_item(
'Invalid request... sending PoM 2')
await self.send_pom(Misbehaviour.BAD_PROMISE_ACCEPT, message,
connection)
return
Exchange.add_signature(message.token.bn_signature, message.signature)
encrypted_promised_txs = self.encrypt_txs(message.type, message.needed,
message.promised,
message.token.epoch)
briefcase_message = BriefcaseBARMessage(message.token, message.to_peer,
message.from_peer, message,
encrypted_promised_txs)
briefcase_message.set_byzantine(self.config.get('byzantine'))
briefcase_message.compute_signature()
await self.send(connection, briefcase_message)
def im_included(self, peer_list):
for peer in peer_list:
if peer.public_address == '{}:{}'.format(self.config.get('host'),
self.config.get('port')):
Logger.get_instance().debug_item(
'I am included in this epoch!')
return True
Logger.get_instance().debug_item('I am excluded in this epoch!')
return False
def set_new_epoch(cls):
current_epoch: Epoch = cls.get_current_epoch()
cls.update_to_old(current_epoch)
next_epoch = cls.get_next_epoch()
cls.update_to_current(next_epoch)
cls.add(Epoch(epoch=next_epoch.epoch + 1, current=False, next=True))
Logger.get_instance().debug_item(
'Current epoch: {}, next epoch: {}.'.format(
next_epoch.epoch, next_epoch.epoch + 1))
return next_epoch.epoch
async def freeze(self, epoch):
async with self.lock:
if self._frozen_mp_epoch != epoch:
self._frozen_mp_epoch = epoch
self.frozen_mp = set(self.mp)
Logger.get_instance().debug_item('Mempool frozen!',
LogLevels.INFO)
else:
Logger.get_instance().debug_item('Mempool already frozen!',
LogLevels.INFO)
async def drop_connection(self, connection: StreamWriter) -> NoReturn:
address = self.to_address(connection.get_extra_info('peername'))
if address not in self.connections: return
conn = self.connections[address]
if not conn.is_closing():
connection.close()
await connection.wait_closed()
self.connections.pop(address)
Logger.get_instance().debug_item(
'Connection with {} closed'.format(address))
def insert(self, txs: List[MempoolDisk], added: List[str]):
old_size = self.size
for tx in txs:
if tx.short_id in added:
self.mp.add(tx.short_id)
self._mapping[tx.short_id] = tx.full_id
self.size = self.size + 1
Logger.get_instance().debug_item(
'I had {} txs and now I have {} txs in my mempool'.format(
old_size, self.size), LogLevels.INFO)
async def on_message(self, message: Message):
if message.author.bot or message.channel.is_nsfw():
return
Logger.debug(f"Started processing {message.id}")
new_images = []
if message.attachments:
process = []
for attachment in message.attachments:
for extension in image_extensions:
if attachment.filename.endswith(extension):
process.append(attachment)
for attachment in process:
await attachment.save(f"./tmp/{attachment.filename}")
classification = classifier.classify(
f"./tmp/{attachment.filename}"
)[f"./tmp/{attachment.filename}"]
if classification["safe"] < classification["unsafe"]:
name = f"./tmp/{round(time() * 1000)}-{attachment.filename}"
detector.censor(f"./tmp/{attachment.filename}", name)
new_images.append(File(name, attachment.filename))
if new_images:
hooks = await message.channel.webhooks()
sender = None
for hook in hooks:
if hook.user == self.bot.user:
sender = hook
break
if sender is None:
sender = await message.channel.create_webhook(
name="Guardian",
reason="No existing webhook was found, creating one!")
for image in new_images:
em = await self.embed(
sender,
"",
get_embed=True,
image=f"attachment://{image.filename}",
footer=Footer("Auto detected NSFW, removed by Guardian",
self.bot.user.avatar_url))
# await self.send(sender, "", embed=em, file=image)
await sender.send(content=message.content,
username=str(message.author),
avatar_url=message.author.avatar_url,
file=image,
embed=em)
await message.delete()
Logger.debug(f"Successfully processed {message.id}")
def decrypt_briefcase(self, briefcase, key):
try:
briefcase = json.loads(briefcase)
if briefcase is None:
return None
for index, encrypted_tx in enumerate(briefcase):
briefcase[index] = self.crypto.get_aes().decrypt(encrypted_tx.encode(), key).decode()
return [Data(bytes.fromhex(tx)) for tx in briefcase]
except Exception as e:
Logger.get_instance().debug_item('Invalid briefcase decrypt, {}'.format(e))
return None
def __init__(self):
super().__init__(Intents.all())
self.prefix = self.get_bot_prefix
self.description = "Guardian protects your server against NSFW!"
modules = list(
map(lambda extension: extension.replace("/", ".")[:-3],
glob("src/modules/*.py")))
for index, _ in enumerate(self.load_extensions(modules)):
Logger.info(
f"Loaded: {modules[index].replace('src.modules.', '')}")
async def register_to_bn(self) -> NoReturn:
msg = HelloMessage(self.public_key, self.host, self.port)
for bn in self.bootstrap_nodes_addresses:
bn = self.set_bn(bn)
bn_ip, bn_port = self.get_ip_port(bn.address)
BootstrapIdentity.get_or_add(bn)
Logger.get_instance().debug_item(
'Sending Hello message to: {}:{}'.format(bn_ip, bn_port))
try:
await self.send_to(bn_ip, bn_port, msg)
except Exception as e:
continue
async def _handle(self, connection: StreamWriter, message: ConnectionRequestBARMessage) -> NoReturn:
if not await self.is_valid_message(message):
Logger.get_instance().debug_item('Invalid seed ... sending PoM')
await self.send_pom(Misbehaviour.BAD_SEED, message, connection)
else:
mempool_dump = await self.mempool.serialize()
history_divulge_message = HistoryDivulgeBARMessage(mempool_dump, message.token,
message.to_peer, message.from_peer, message)
history_divulge_message.set_byzantine(self.config.get('byzantine'))
history_divulge_message.compute_signature()
await self.send(connection, history_divulge_message)
def add_if_new(cls, txs):
added = []
for tx in txs:
exists = cls.get_session().query(cls).filter(
or_(cls.full_id == tx.full_id,
cls.short_id == tx.short_id)).first()
if not exists:
added.append(tx.short_id)
cls.add(tx)
else:
Logger.get_instance().debug_item(
'duplicate {}'.format(tx.short_id), LogLevels.INFO)
return added
async def handle(self, msg_bytes: bytes,
connection: StreamWriter) -> NoReturn:
message = self.deserialize_data(msg_bytes)
for controller in self.controllers:
if controller.is_valid_controller_for(message):
Logger.get_instance().debug_item(
'A {} message is arrived and being handled'.format(
message))
try:
await controller.handle(connection, message)
except Exception as _:
break
finally:
break
async def _handle(self, connection: StreamWriter,
message: HistoryDivulgeBARMessage) -> NoReturn:
if not await self.is_valid_message(message):
await self.send_pom(Misbehaviour.BAD_SEED, message, connection)
Logger.get_instance().debug_item('Invalid request ... sending PoM')
else:
partner_mempool = Mempool.deserialize(message.elements)
intersection_set_a_b = Mempool.get_diff(self.mempool.frozen_mp,
partner_mempool)
intersection_set_b_a = Mempool.get_diff(partner_mempool,
self.mempool.frozen_mp)
exchange_type, exchange_number = self.bal_or_opt_exchange(
intersection_set_b_a, intersection_set_a_b)
if exchange_type == self.ABORT:
exchange = Exchange(seed=message.token.bn_signature,
sender=True,
needed=json.dumps([]),
promised=json.dumps([]),
type=exchange_type,
signature='',
valid=True)
Exchange.add(exchange)
return
needed, promised = await self.select_exchanges(
exchange_type, intersection_set_b_a, intersection_set_a_b,
exchange_number)
ser_needed, ser_promised = json.dumps(needed), json.dumps(promised)
exchange = Exchange(seed=message.token.bn_signature,
sender=True,
needed=ser_needed,
promised=ser_promised,
type=exchange_type,
signature='',
valid=False)
Exchange.add(exchange)
exchange_message = ExchangeBARMessage(message.token,
message.to_peer,
message.from_peer, message,
needed, promised,
exchange_type)
exchange_message.set_byzantine(self.config.get('byzantine'))
exchange_message.compute_signature()
await self.send(connection, exchange_message)
def is_promise_request_valid(self, message: ExchangeBARMessage):
for tx in message.needed:
if not self.mempool.has(tx):
Logger.get_instance().debug_item('Wrong needed',
LogLevels.ERROR)
return False
for tx in message.promised:
if self.mempool.has(tx):
Logger.get_instance().debug_item('Wrong promised',
LogLevels.ERROR)
return False
if not self.is_bar_or_opt(message) and not self.is_valid_bar(
message) and not self.is_valid_opt(message):
return False
return True
def start(self) -> NoReturn:
Logger.get_instance().debug_item(
'Starting {} node {} on port {}'.format(self.__class__.__name__,
self.id, self.port),
LogLevels.PRODUCTION)
try:
self.loop.run_forever()
except KeyboardInterrupt:
Logger.get_instance().debug_item('KeyboardInterrupt, exiting...',
LogLevels.ERROR)
for task in asyncio.Task.all_tasks():
task.cancel()
self.server.close()
self.loop.run_until_complete(self.server.wait_closed())
self.loop.close()
async def _handle(self, connection: StreamWriter, message: PoMBARMessage):
peer_from_pk = self.crypto.get_ec().public_key_to_string(
message.from_peer.public_key)
peer_to_pk = self.crypto.get_ec().public_key_to_string(
message.to_peer.public_key)
peer_from = Peer.find_one_by_public_key(peer_from_pk)
peer_to = Peer.find_one_by_public_key(peer_to_pk)
pom = ProofOfMisbehaviour(against_peer=peer_from.id,
type=message.misbehaviour.value,
from_peer=peer_to.id,
epoch=message.token.epoch)
Logger.get_instance().debug_item(
'Received a PoM against peer: {} from peer: {}'.format(
peer_from.public_address, peer_to.public_address))
ProofOfMisbehaviour.add(pom)
async def _handle(self, connection: StreamWriter,
message: KeyRequestBARMessage):
if not await self.is_valid_message(message):
Logger.get_instance().debug_item('Invalid request... sending PoM')
await self.send_pom(Misbehaviour.BAD_SEED, message, connection)
else:
key = Token.find_one_by_epoch(message.token.epoch).key
key_message = KeyBARMessage(message.token, message.to_peer,
message.from_peer, message, key)
key_message.set_byzantine(self.config.get('byzantine'))
key_message.compute_signature()
await self.send(connection, key_message)
def test_ckps(data_dir,
auged,
ckp_dir,
log_dir,
model,
model_kwargs,
mets,
device=torch.device("cpu"),
loss_func=torch.nn.BCELoss(),
total_amt=16384,
val_percent=0.25,
test_amt=768,
wrapped_function=None,
workers=0,
seed=42):
"""
:param data_dir: the directory to the data
:param auged: whether or not the data is augmented or not
:param ckp_dir: the directory to the checkpoint
:param log_dir: the directory to the logs
:param model: the class of the model
:param mets: the metrics to use
:param device: the device to use
:param loss_func: the loss function to use
:param total_amt: the total amount of data to use
:param val_percent: the percent of data to use for validation
:param test_amt: the amount of data to use for testing
:param wrapped_function: the wrapped function ot use
:param workers: the number of workers to use
:param seed: the seed to use
"""
data = Data(data_dir,
auged,
total_amt=total_amt,
val_percent=val_percent,
test_amt=test_amt,
wrapped_function=wrapped_function,
workers=workers,
device=device,
verbose=True,
seed=seed)
test_data = data.get_test_data()
for dir_ in os.listdir(ckp_dir):
mod_ckp_dir = f"{ckp_dir}/{dir_}"
if os.path.isdir(mod_ckp_dir):
for ckp in os.listdir(mod_ckp_dir):
if "FINAL" in ckp:
fin_dir = f"{mod_ckp_dir}/{ckp}"
name = ckp.replace(".pt", "")
logger = Logger(f"{name}_TEST",
log_dir,
mets,
overwrite=True,
verbose=True)
mod = model(**model_kwargs)
mod.to(device)
print(f"Loading model from {fin_dir}")
mod, _ = load_ckp(fin_dir, mod, dev=device)
name = name.replace("_FINAL.pt", "")
test_model(test_data, mod, loss_func, logger, name)
async def _handle(self, connection: StreamWriter,
message: ExchangeBARMessage) -> NoReturn:
if not await self.is_valid_message(
message) or not self.is_promise_request_valid(message):
Logger.get_instance().debug_item('Invalid request... sending PoM')
await self.send_pom(Misbehaviour.BAD_SEED, message, connection)
# if not self.is_promise_request_valid(message):
# Logger.get_instance().debug_item('Invalid history message... sending PoM')
# await self.send_pom(Misbehaviour.BAD_PROMISE, message, connection)
# return
else:
ser_needed, ser_promised = json.dumps(message.needed), json.dumps(
message.promised)
exchange = Exchange(seed=message.token.bn_signature,
sender=False,
needed=ser_promised,
promised=ser_needed,
type=str(message.type),
signature=message.signature,
valid=False)
Exchange.add(exchange)
promise_message = PromiseBARMessage(message.token, message.to_peer,
message.from_peer, message,
message.promised,
message.needed,
str(message.type))
promise_message.set_byzantine(self.config.get('byzantine'))
promise_message.compute_signature()
encrypted_promised_txs = self.encrypt_txs(message.type,
message.needed,
message.promised,
message.token.epoch)
briefcase_message = BriefcaseBARMessage(message.token,
message.to_peer,
message.from_peer, message,
encrypted_promised_txs)
briefcase_message.set_byzantine(self.config.get('byzantine'))
briefcase_message.compute_signature()
await self.send(connection, promise_message)
await self.send(connection, briefcase_message)
async def _handle(self, connection: StreamWriter, message: RenewTokenMessage):
is_valid_token = message.is_valid_signature()
current_view_peers = ViewMessage.get_current_view()
was_peer_honest = self.was_honest_peer(message, current_view_peers)
if is_valid_token:
current_view_peers = ViewMessage.get_current_view()
view_message = ViewMessage(peer_list=current_view_peers, epoch=self.get_current_epoch().epoch)
Logger.get_instance().debug_list(view_message.peer_list, separator='\n')
token_message = self.create_token(message.base, message.proof)
view_message.set_token(token_message)
await self.send(connection, view_message)
peer_address = self.format_address(connection.get_extra_info('peername'))
peer_pk = self.get_public_key(message.base)
peer = Peer.find_on_by_address_or_pk(peer_address, peer_pk)
current_epoch = self.get_current_epoch()
View.add(View(peer=peer.id, epoch_id=current_epoch.id))
Logger.get_instance().debug_item('Renewed View Message for epoch {} sent!'.format(current_epoch.epoch))
def init(self):
txs = MempoolDisk.get_all()
for tx in txs:
self.mp.add(tx.short_id)
self._mapping[tx.short_id] = tx.full_id
for i in range(MAX_FAKE_DATA):
transaction = bytearray(
random.getrandbits(8) for _ in range(
int(
random.normalvariate(TRANSACTION_SIZE_MU,
TRANSACTION_SIZE_SIGMA))))
self.fake_data.add(transaction.hex())
self.size = len(txs)
self.frozen_mp = set(self.mp)
Logger.get_instance().debug_item(
'I have {} txs in my mempool'.format(len(txs)), LogLevels.INFO)
return len(txs) == len(self.mp)
async def _handle(self, connection: StreamWriter, message: HelloMessage):
difficulty = self.get_puzzle_difficulty()
current_epoch = self.get_current_epoch().epoch
register_message = RegisterMessage(difficulty, message.public_key,
current_epoch)
already_exist = Registration.is_registration_present(
register_message.puzzle)
if already_exist:
register_message.puzzle = already_exist.base
Logger.get_instance().debug_item(
'A valid registration already exist')
await self.send(connection, register_message)
else:
epoch = self.get_current_epoch().epoch
Registration.add(
Registration(base=register_message.puzzle, epoch=epoch))
await self.send(connection, register_message)
async def _handle(self, connection: StreamWriter, message: KeyBARMessage):
if not await self.is_valid_message(message):
Logger.get_instance().debug_item('Invalid request... sending PoM')
await self.send_pom(Misbehaviour.BAD_SEED, message, connection)
else:
exchange = Exchange.get_exchange(message.token.bn_signature)
data = self.decrypt_briefcase(exchange.briefcase, message.key)
if not self.is_valid_data(exchange.needed, data):
Logger.get_instance().debug_item('Invalid data... sending PoM')
await self.send_pom(Misbehaviour.BAD_BRIEFCASE, message, connection)
return
Exchange.set_valid(message.token.bn_signature)
txs = [MempoolDisk(data=tx.data, short_id=tx.short_hash, full_id=tx.hash) for tx in data]
real_txs = list(filter(lambda tx: tx.short_id in set(json.loads(exchange.needed)), txs))
added = MempoolDisk.add_if_new(real_txs)
self.mempool.insert(txs, added)
def activate_logger(self):
if not self.__check_logger_status():
GS.LOGGER = Logger()
GS.LOGGER.write_to_logger(
"########################################################################\n"
"########################### LOGGER ACTIVATED ###########################\n"
"########################################################################"
)
return GS.LOGGER
async def start_new_epoch(self):
while True:
key, view_message = await PubSub.get_subscriber_epoch_instance(
).consume()
next_epoch_time = int(view_message.next_epoch) - int(
datetime.now(tz=timezone.utc).timestamp())
await asyncio.sleep(next_epoch_time)
renew_message = RenewTokenMessage(view_message.token.base,
view_message.token.proof,
view_message.token.bn_signature,
view_message.token.epoch)
Logger.get_instance().debug_item(
'Number of open connections: {}'.format(
len(self.connections.keys())))
bn = BootstrapIdentity.get_one_by_token(renew_message.bn_signature)
await Mempool().get_instance().freeze(view_message.epoch)
writer = self.connections[bn.address]
writer.write(renew_message.serialize())
Logger.get_instance().debug_item('Renew token message sent!')
await writer.drain()
def deserialize_data(self, msg_bytes) -> Message:
try:
return None if msg_bytes == b'' or msg_bytes is None else Message.deserialize(
msg_bytes)
except Exception as _:
Logger.get_instance().debug_item(
'deserialize exception: {}'.format(msg_bytes), LogLevels.ERROR)
self.invalid_messages.append(msg_bytes)
if len(self.invalid_messages) > 1:
try:
print('start')
t = b''.join(self.invalid_messages)
print('t done', t)
msg = Message.deserialize(t)
print('recovery with {} messages',
len(self.invalid_messages))
self.invalid_messages = []
Logger.get_instance().debug_item(
'message recovered successfully', LogLevels.INFO)
return msg
except Exception as e:
self.invalid_messages.append(msg_bytes)
Logger.get_instance().debug_item(
'failed message recovery: {}'.format(e),
LogLevels.ERROR)
async def _handle(self, connection: StreamWriter,
message: RegisterMessage):
bn_address = self.format_address(connection.get_extra_info('peername'))
already_registered: Token = Token.find_one_by_address(
bn_address, message.current_epoch)
if already_registered:
Logger.get_instance().debug_item('Found a valid token!',
LogLevels.INFO)
login_message = LoginMessage(already_registered.base,
already_registered.proof,
self.config.get_address())
await self.send(connection, login_message)
else:
Logger.get_instance().debug_item('Computing a valid PoW ...',
LogLevels.INFO)
pow_solution = Hashcash.new(message.difficulty,
message.puzzle.encode('utf-8'))
Logger.get_instance().debug_item(
'PoW found! Salt: {}, percentile: {}'.format(
pow_solution.salt.hex(), pow_solution.percentile()),
LogLevels.INFO)
login_message = LoginMessage(message.puzzle,
pow_solution.salt.hex(),
self.config.get_address())
await self.send(connection, login_message)
