154 lines
4.4 KiB
Python
154 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import kitty
|
|
|
|
import pyvista as pv
|
|
import numpy as np
|
|
import math
|
|
import subprocess
|
|
import re
|
|
import requests
|
|
import argparse
|
|
from io import BytesIO
|
|
|
|
# TODO: Color arches based on latency
|
|
# TODO: Interactive globe (spin w/ keys)
|
|
# TODO: Text info (num hops etc.)
|
|
# TODO: Image spacing
|
|
# TODO: Async rendering?
|
|
|
|
|
|
# Convert lat/lon to Cartesian coordinates
|
|
def latlon_to_xyz(lat: float, lon: float, radius=1.0):
|
|
lat_rad = np.radians(lat)
|
|
lon_rad = np.radians(lon)
|
|
x = radius * np.cos(lat_rad) * np.cos(lon_rad)
|
|
y = radius * np.cos(lat_rad) * np.sin(lon_rad)
|
|
z = radius * np.sin(lat_rad)
|
|
return np.array([x, y, z])
|
|
|
|
|
|
# Create an arch between two 3D points
|
|
def generate_arch(p1, p2, height_factor=0.2, n_points=100):
|
|
# Normalize input points to lie on the unit sphere
|
|
p1 = p1 / np.linalg.norm(p1)
|
|
p2 = p2 / np.linalg.norm(p2)
|
|
|
|
# Compute angle between p1 and p2
|
|
omega = np.arccos(np.clip(np.dot(p1, p2), -1, 1))
|
|
if omega == 0:
|
|
return np.tile(p1, (n_points, 1)) # degenerate case
|
|
|
|
t = np.linspace(0, 1, n_points)
|
|
sin_omega = np.sin(omega)
|
|
|
|
# Spherical linear interpolation (slerp)
|
|
arch_points = (
|
|
np.sin((1 - t)[:, None] * omega) * p1[None, :]
|
|
+ np.sin(t[:, None] * omega) * p2[None, :]
|
|
) / sin_omega
|
|
|
|
# Add radial height offset based on sine curve
|
|
heights = 1 + np.sin(np.pi * t) * height_factor
|
|
arch_points *= heights[:, None] # Scale outward from center
|
|
|
|
return arch_points
|
|
|
|
|
|
def traceroute(target: str) -> list[tuple[int, int]]:
|
|
# Run traceroute command
|
|
result = subprocess.run(
|
|
["traceroute", "-n", target, "-q", "1", "-w", "1,3,10"],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
hops: list[str] = re.findall(r"\n\s*\d+\s+([\d.]+)", result.stdout)
|
|
|
|
coords: list[tuple[int, int]] = []
|
|
for ip in hops:
|
|
try:
|
|
response: dict = requests.get(f"http://ip-api.com/json/{ip}").json()
|
|
if response["status"] == "success":
|
|
coords.append((response["lat"], response["lon"]))
|
|
except Exception:
|
|
continue
|
|
return coords
|
|
|
|
|
|
def main():
|
|
|
|
parser = argparse.ArgumentParser(
|
|
prog="kglobe",
|
|
description="Traceroute on a globe",
|
|
epilog="Requires kitty graphics protocol support in terminal",
|
|
)
|
|
|
|
parser.add_argument("-t", "--traceroute", default=None)
|
|
parser.add_argument("-e", "--external", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
|
|
locations = []
|
|
if args.traceroute:
|
|
locations = traceroute(args.traceroute)
|
|
|
|
globe = pv.Sphere(
|
|
radius=1.0,
|
|
theta_resolution=120,
|
|
phi_resolution=120,
|
|
start_theta=270.001,
|
|
end_theta=270,
|
|
)
|
|
|
|
tex = pv.examples.load_globe_texture()
|
|
|
|
globe.active_texture_coordinates = np.zeros((globe.points.shape[0], 2))
|
|
|
|
globe.active_texture_coordinates[:, 0] = 0.5 + np.arctan2(
|
|
globe.points[:, 1], globe.points[:, 0]
|
|
) / (2 * math.pi)
|
|
globe.active_texture_coordinates[:, 1] = (
|
|
0.5 + np.arcsin(globe.points[:, 2]) / math.pi
|
|
)
|
|
|
|
# Convert to 3D coordinates
|
|
points_3d = [latlon_to_xyz(lat, lon) for lat, lon in locations]
|
|
|
|
pl: pv.Plotter = pv.Plotter(off_screen=(not args.external))
|
|
pl.add_mesh(globe, color="tan", smooth_shading=True, texture=tex, show_edges=False)
|
|
|
|
for pt in points_3d:
|
|
city_marker = pv.Sphere(center=pt, radius=0.02)
|
|
pl.add_mesh(city_marker, color="blue")
|
|
for i in range(len(points_3d[:-1])):
|
|
arch = generate_arch(points_3d[i], points_3d[i + 1], height_factor=0.2)
|
|
line = pv.lines_from_points(arch, close=False)
|
|
pl.add_mesh(line, color="red", line_width=2)
|
|
|
|
kitty.hide_cursor()
|
|
y, x = kitty.get_position()
|
|
|
|
height, width = kitty.get_terminal_size_pixel()
|
|
h_pix, w_pix = kitty.get_terminal_cell_size()
|
|
|
|
# the image requires height/cell_height lines
|
|
needed_lines = round(height/h_pix)
|
|
print("\n" * needed_lines, end="")
|
|
|
|
try:
|
|
if not args.external:
|
|
while True:
|
|
pl.camera.Azimuth(1)
|
|
buf: BytesIO = BytesIO()
|
|
pl.screenshot(buf, transparent_background=True, window_size=(height, width))
|
|
kitty.set_position(y - needed_lines, 0)
|
|
kitty.draw_to_terminal(buf)
|
|
else:
|
|
pl.show()
|
|
finally:
|
|
kitty.show_cursor()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|