def to_device(self):
# Torch CUDA config
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if not torch.cuda.is_available():
printc("GPU NOT AVAILABLE, USING CPU!", color="ORANGE")
self.model.to(self.device)
cudnn.benchmark = True # For fast training.
def init_server(port, K=[], _tencrypt=None, _tdecrypt=None, _getdate=None, _dest='', _destport=0, _filt=''):
global KEYS, tencrypt, tdecrypt, filt, dest, destport, gport, getdate
gport = port
KEYS = K
tencrypt = _tencrypt
tdecrypt = _tdecrypt
getdate = _getdate
filt = _filt
dest = _dest
destport = _destport
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(("localhost", 80))
printc("Established server on %s:%d" % ("localhost", 80), style.bold, style.blue)
sock.listen(1)
# Message port
msock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
msock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
host = '0.0.0.0'
hport = port
msock.bind((host, hport))
printc("Filtered receive port established on %s:%d" % (host, hport), style.bold, style.blue)
msock.listen(1)
wait_for_connections([sock, msock])
def save_metrics(self):
self.metrics = self.pruning_metrics()
with open(self.path / 'metrics.json', 'w') as f:
json.dump(self.metrics, f, indent=4)
printc(json.dumps(self.metrics, indent=4), color='GRASS')
summary = self.pruning.summary()
summary_path = self.path / 'masks_summary.csv'
summary.to_csv(summary_path)
print(summary)
def load_keys(filename, verbose=False):
global KEYS
if verbose:
printc("Loading one-time-pads from %s" % (os.path.abspath(filename)),
style.blue)
with open(os.path.abspath(filename), "r") as obj:
KEYS = obj.read().split("\n")
printc("Successfully loaded keys", style.bold)
def run(self):
self.freeze()
printc(f"Running {repr(self)}", color='YELLOW')
self.to_device()
self.build_logging(self.train_metrics, self.path)
self.save_metrics()
if self.pruning.compression > 1:
self.run_epochs()
def generate(K, N, filename, filt='', verbose=False):
apath = os.path.abspath(filename)
if verbose:
printc("Writing %d keys of length %d to %s" % (K, N, apath),
style.blue)
r = [gen_str(N, filt=filt) for i in range(K)]
with open(apath, "w+") as obj:
obj.write('\n'.join(r))
def _send_message(sock, m, filt='', verbose=False):
try:
sock.send(m)
write_message(('localhost', 0), m, filt=filt)
if verbose:
printc("Message sent successfully", style.green)
except:
if verbose:
printc("Error: Message failed to send. Retrying...",
style.highlightred)
return _send_message(sock, m, verbose=verbose)
def _receive_message(sock, verbose=False):
sock.listen(5)
while True:
conn, addr = sock.accept()
if verbose:
printc(
"Client %s:%d has connected and may or may not be able to read and send messages"
% (addr[0], addr[1]), style.blue)
start_new_thread(client_thread, (conn, addr))
conn.close()
def client_thread(conn, addr):
buffer = []
while True:
data = conn.recv(8192)
if not data:
break
write_message(addr, data)
printc("Client [%s:%d] has dropped their connection" % (addr[0], addr[1]),
style.red)
conn.close()
def do_chat(addr, dport, filt='', verbose=False):
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
if verbose:
printc("Connecting to %s:%d ..." % (addr, dport), style.blue)
sock.connect((addr, dport))
if verbose:
printc("Successfully connected to %s:%d" % (addr, dport), style.bold,
style.green)
while True:
m = raw_input(style.bold + "Message: \033[0m")
_send_message(sock, m, filt=filt, verbose=verbose)
def build_logging(self, metrics, path=None, csv=True, tensorboard=False):
if path is None:
self.path = self.get_path()
printc(f"Logging results to {self.path}", color='MAGENTA')
print('Type-path::', type(self.path))
if isinstance(self.path, str):
self.path = pathlib.Path(self.path)
self.path.mkdir(exist_ok=True, parents=True)
self.save_params()
self.log_csv = csv
self.log_tb = tensorboard
self.log_epoch_n = 0
if self.log_csv:
self.csvlogger = CSVLogger(self.path / 'logs.csv', metrics)
if self.log_tb:
tb_path = self.path / 'tbevents'
tb_path.mkdir()
self.tblogger = SummaryWriter(log_dir=tb_path)
def run_epochs(self):
since = time.time()
try:
for epoch in range(self.epochs):
printc(f"Start epoch {epoch}", color='YELLOW')
self.train(epoch)
self.eval(epoch)
# Checkpoint epochs
# TODO Model checkpointing based on best val loss/acc
if epoch % self.save_freq == 0:
self.checkpoint()
# TODO Early stopping
# TODO ReduceLR on plateau?
self.log(timestamp=time.time()-since)
self.log_epoch(epoch)
except KeyboardInterrupt:
printc(f"\nInterrupted at epoch {epoch}. Tearing Down", color='RED')
def modify(filename, filt='', verbose=False):
global CHRS
apath = os.path.abspath(filename)
if verbose:
printc("Modifying %s with new filter ..." % apath, style.blue)
if filt == '':
return
lines = []
cl, fl = len(CHRS), len(filt)
if verbose:
printc("Sanitizing..." % apath, style.blue)
with open(apath, "r") as obj:
lines = obj.read().split("\n")
for i, line in enumerate(lines):
for j, c in enumerate(line):
lines[i][j] = (CHRS.index(c) + CHRS.index(filt[j % fl])) % cl
if verbose:
printc("Writing new sanitization to file..." % apath, style.blue)
with open(apath, "w+") as obj:
obj.write('\n'.join(lines))
def log_chat(sock, verbose=False, filt=''):
last_count = 0
printc("Listening...", style.bold)
while True:
if not os.path.exists("log"):
with open("log", "w+") as obj:
obj.write("")
with open("log", "r") as obj:
lines = obj.read().split("\n")
llen = len(lines)
if llen != last_count:
for i in range(last_count, llen):
if lines[i].replace(' ', '') == '':
continue
log_message(tdecrypt(lines[i], filt=filt))
last_count = llen
time.sleep(3)
def drawBoard(n, labels, assignments, bridges):
for i in range(n):
for j in range(n):
if (i == 0 or i == n - 1) and (j == 0 or j == n - 1):
print "%-*s" % (width, ""),
else:
if (i, j) in labels:
if labels[(i, j)] in assignments[0]:
#printc("%-*s" % (width, assignments[0][labels[(i,j)]])),
if (i, j) in bridges:
printc('O', bcolors.OKBLUE, 'bold'),
else:
printc('O'),
elif labels[(i, j)] in assignments[1]:
#printc("%-*s" % (width, assignments[1][labels[(i,j)]]), bcolors.FAIL),
if (i, j) in bridges:
printc('O', bcolors.FAIL, 'bold'),
else:
printc('O', bcolors.FAIL),
else:
print "%-*s" % (width, "x"),
if j == n - 1:
print "\n"
def pprint(self):
width = 4
for i in range(self.N):
for j in range(self.N):
if (i == 0 or i == self.N - 1) and (j == 0 or j == self.N - 1):
print "%-*s" % (width, ""),
else:
if (i, j) in self.label:
if self.label[(i, j)] in self.assignment[0]:
#printc("%-*s" % (width, assignments[0][labels[(i,j)]])),
if (i, j) in self.bridges:
printc('O', bcolors.OKBLUE, 'bold'),
else:
printc('O'),
elif self.label[(i, j)] in self.assignment[1]:
#printc("%-*s" % (width, assignments[1][labels[(i,j)]]), bcolors.FAIL),
if (i, j) in self.bridges:
printc('O', bcolors.FAIL, 'bold'),
else:
printc('O', bcolors.FAIL),
else:
print "%-*s" % (width, "x"),
if j == self.N - 1:
print "\n"
def dolisten(port, filt='', verbose=False):
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
if verbose:
printc("Attempting to bind server socket to localhost...", style.blue)
try:
sock.bind(('', port))
except:
printc(
"Error: Could not bind server socket to localhost. Perhaps a zombie process is already running?",
style.bold, style.highlightred)
exit(1)
if verbose:
printc("Successfully bound socket to localhost...", style.green)
start_new_thread(log_chat, (sock, verbose, filt))
_receive_message(sock, verbose=verbose)
def main():
global KEYS
printc("\n ~~whisper a.3", style.bold)
printc(" silent communication\n", style.bold, style.purple)
printc(
"Legal Disclaimer: It is the end user's responsibility to obey all local, state, and federal laws. The developer assumes no liability for any misuse or damages caused by this software.\n\n",
style.yellow)
listen = False
verbose = False
serve = False
their_addr = ''
port = 9504
destport = 9504
keyfile = 'whisp'
gen_file = ''
gen_num = 100000
gen_len = 1024
mod_file = ''
filt = ''
out_filt = ''
i = 1
if len(sys.argv) < 2:
help()
return
while i < len(sys.argv):
o = sys.argv[i]
if o == "-h" or o == "--help":
help()
return
elif o == "--listen":
listen = True
elif o == "--serve":
serve = True
elif o == "-s":
filt = sys.argv[i + 1]
i += 1
elif o == '-d':
their_addr = sys.argv[i + 1]
i += 1
elif o == '-w':
keyfile = sys.argv[i + 1]
i += 1
elif o == '-p':
port = int(sys.argv[i + 1])
i += 1
elif o == '-dp':
destport = int(sys.argv[i + 1])
i += 1
elif o == '-g':
gen_file = sys.argv[i + 1]
i += 1
elif o == '-m':
gen_file = sys.argv[i + 1]
i += 1
elif o == '-f':
out_filt = sys.argv[i + 1]
i += 1
elif o == '-k':
gen_num = int(sys.argv[i + 1])
i += 1
elif o == '-l':
gen_len = int(sys.argv[i + 1])
i += 1
elif o == '-v':
verbose = True
else:
printc("Warning: Ignoring unknown parameter %s" % o, style.yellow)
i += 1
if gen_file != '':
generate(gen_num, gen_len, gen_file, filt=out_filt, verbose=verbose)
elif mod_file != '':
modify(mod_file, filt=out_filt, verbose=verbose)
load_keys(keyfile, verbose=verbose)
if serve:
init_server(port, KEYS, tencrypt, tdecrypt, getdate, their_addr,
destport, filt)
elif listen:
dolisten(port, filt=filt, verbose=verbose)
else:
do_chat(their_addr, destport, filt=filt, verbose=verbose)
def apply_pruning(self, strategy, compression):
constructor = getattr(strategies, strategy)
x, y = next(iter(self.train_dl))
self.pruning = constructor(self.model, x, y, compression=compression)
self.pruning.apply()
printc("Masked model", color='GREEN')
