71 lines
1.9 KiB
Python
71 lines
1.9 KiB
Python
import numpy as np
|
|
from pwnlib import tubes
|
|
import time
|
|
import matplotlib.pyplot as plt
|
|
from rmsd import calculate_rmsd
|
|
from scipy.spatial.transform import Rotation
|
|
import seaborn as sb
|
|
import itertools
|
|
from sklearn.neighbors import BallTree
|
|
|
|
def read_starfile(data):
|
|
stars = []
|
|
for line in data.strip().split('\n'):
|
|
[x,y,z,m] = [float(s.strip()) for s in line.split(',')]
|
|
stars.append({'v': np.array([x,y,z]), 'm':m})
|
|
return stars
|
|
|
|
with open('./myspace-mike33686zulu/test.txt') as f:
|
|
catalog = read_starfile(f.read())
|
|
|
|
def angle(x, y):
|
|
x_ = x/np.linalg.norm(x)
|
|
y_ = y/np.linalg.norm(y)
|
|
return np.arccos(np.clip(np.dot(x_, y_), -1.0, 1.0))
|
|
|
|
def gen_fingerprints(X, bt):
|
|
dist, ind = bt.query(X[:], k=3)
|
|
fingerprints = []
|
|
for i in range(X.shape[0]):
|
|
a = X[i,:]
|
|
[bi, ci] = ind[i,1:]
|
|
b = X[bi,:]
|
|
c = X[ci,:]
|
|
fingerprints.append([
|
|
*dist[i, 1:],
|
|
angle(b-a, c-a),
|
|
])
|
|
return fingerprints
|
|
|
|
TICKET = 'THE_TICKET'
|
|
r = tubes.remote.remote('myspace.satellitesabove.me', 5016)
|
|
r.send(TICKET+'\n')
|
|
time.sleep(0.5)
|
|
r.recvuntil('Ticket please:\n', drop=True)
|
|
|
|
X = np.vstack([s['v'] for s in catalog])
|
|
bt_catalog = BallTree(X, leaf_size=30)
|
|
|
|
fingerprints = gen_fingerprints(X, bt_catalog)
|
|
|
|
for _ in range(5):
|
|
refstars = read_starfile(r.recvuntil('\n\n').decode())
|
|
|
|
Y = np.vstack([s['v'] for s in refstars])
|
|
bt_unknown = BallTree(Y, leaf_size=30)
|
|
|
|
fingerprints_unknown = gen_fingerprints(Y, bt_unknown)
|
|
|
|
matches = []
|
|
for query in fingerprints_unknown:
|
|
error = [np.linalg.norm(np.array(fp_known) - np.array(query))
|
|
for fp_known in fingerprints]
|
|
match_idx = np.argmin(error)
|
|
if error[match_idx] < 1e-4:
|
|
print(match_idx, error[match_idx])
|
|
matches.append(match_idx)
|
|
|
|
r.send(','.join(map(str,matches)) + '\n')
|
|
r.recvuntil('Left...\n')
|
|
print(r.clean())
|