#!/usr/bin/env python3
import socket
import re
import sys
import subprocess
udp_server = ('xmpp.vim-cn.com', 2727)
addr_re = re.compile(r"\('(?P<ip>[^']+)', (?P<port>\d+)(?:, (?P<cport>\d+))?")
def main(server):
  sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  sock.settimeout(2)
  print('Send message...')
  sock.sendto(b'req from holepunch.py\n', udp_server)
  msg, addr = sock.recvfrom(1024)
  print('Got answer from %s%s' % (addr, msg))
  m = addr_re.search(msg.decode())
  if not m:
    print("Error: can't parse answer.")
    sys.exit(1)
  m_ip = m.group('ip')
  m_port = int(m.group('port'))
  port = sock.getsockname()[1]
  print('Got my IP and Port: (%r%s%s).' % (m_ip, m_port, port))
  msg = input('> Peer address: ')
  m = addr_re.search(msg)
  if not m:
    print("Error: can't parse input.")
    sys.exit(2)
  p_ip = m.group('ip')
  p_port = int(m.group('port'))
  c_port = int(m.group('cport'))
  print('send initial packet and wait for answer...')
  sock.sendto(b'HELO\n', (p_ip, p_port))
  try:
    msg = sock.recvfrom(1024)
    print('Received:', msg)
  except socket.timeout:
    print("Timed out (it's normal).")
  if server:
    sock.close()
    print('Starting mosh server...')
    msg = subprocess.check_output(['mosh-server', 'new', '-p', str(port)])
    secret = msg.split()[3].decode()
    print('Connect with:\nMOSH_KEY=%s MOSH_CPORT=%s mosh-client %s %s' % (secret, c_port, m_ip, m_port))
  else:
    print('done.')
if __name__ == '__main__':
  server = len(sys.argv) == 2 and sys.argv[1] == '-s'
  main(server)