#!/usr/bin/python3

#
# chhoyhopper-client.py
#
# Copyright (C) 2021 by University of Southern California
# Written by ASM Rizvi<asmrizvi@usc.edu>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License,
# version 2, as published by the Free Software Foundation.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.
#

import argparse
from datetime import datetime
import dns.resolver
import hashlib
import ipaddress
import math
import os
import socket
import sys
import time

NEXT_ADDRESS_CHANGE = 60 # in second

class Program:
    def __init__(self):
        self.arg = self.parse_args()

    def parse_args(self):
        parser = argparse.ArgumentParser(description = 'This program finds the valid IPv6 address to connect to the moving server.', epilog="""
An IPv6 server will be hopping around over different IPv6 addresses. 
This program will find out the correct, current IPv6 address,
which is a function of the time, shared secret, and salt.
Server and client must share a key before making the connection.

EXAMPLE: An IPv6 client wants to connect to a moving server.

Input: Domain name of the server or the IPv6 address of the server. Other inputs are optional.

Command: 
Connecting to vm18.ant.isi.edu:

        chhoyhopper-client --address vm18.ant.isi.edu

Connecting to 2001:1878:401::8009:1d15:

        chhoyhopper-client --address 2001:1878:401::8009:1d15

Using a different key file (dafault is ./chhoyhopper_key.bin):

        chhoyhopper-client --address 2001:1878:401::8009:1d15 --key /tmp/public.pem

Output: An established connection.

        """)
        
        parser.add_argument('--address', '-a', required=True, help='Enter the service address IPv6 address or domain name.')
        
        parser.add_argument('--keyfile', '-k', default='./chhoyhopper_key.bin', help='Key file shared by the server')
        
        parser.add_argument('--salt', '-s', default='4750', help='Constant salt for generating key')
        
        parser.add_argument('--service', '-e', default='ssh', help='Enter the service name (ssh is default).')
        args = parser.parse_args()

        return args

def check_ipv6(n):
    try:
        socket.inet_pton(socket.AF_INET6, n)
        return True
    except socket.error:
        return False

def sha256GetIp6Address (timeNowMin, key, constant, prefix):

    result = hashlib.sha256()
    result.update(str(timeNowMin).encode())
    result.update(key)
    result.update(str(constant).encode())
    
    hexaDeciNum = prefix

    i = 0
    while i < 13:
        hexaDeciNum = hexaDeciNum + result.hexdigest()[i : i + 4]
        if i < 12:
            hexaDeciNum = hexaDeciNum + ":"
        i = i + 4

    return hexaDeciNum

def resolveDnsAndGettingAddress(domainName):
    try:
        result = dns.resolver.query(domainName, 'AAAA')
        if len(result) > 0:
            IP = result[0]
    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
        exit("Domain name does not exist! Resolving DNS failed for \"" + domainName + "\". Please enter a valid domain name." )
    return IP

def readKey(keyFile):
    if os.path.exists(keyFile):
        f = open(keyFile, 'rb')
        key = f.read()
        f.close()
        return key
    else:
        exit("Key file not found! To connect to the server, you must have the same key file that is being used by the server. This program searches for the key in the /usr/local/var/ directory or the directory you provided with the --keyfile option.")

def main():
    args = Program()
    IP = args.arg.address
    keyFile = args.arg.keyfile
    constant = args.arg.salt
    service = args.arg.service

    if check_ipv6(IP) is False:
        IP = str(resolveDnsAndGettingAddress(IP))

    print("chhoyhopper-client for clear " + IP)
    
    try:
        ipaddress.ip_address(IP)
    except ValueError:
        exit("Please enter a valid IPv6 address or domain.")

    keyVal = readKey(keyFile)

    prefix = str(ipaddress.ip_network(IP).supernet(new_prefix=64))
    prefix = prefix[0:prefix.index('/')]
    #print(prefix)

    timeNowMin = math.floor((time.time()) / NEXT_ADDRESS_CHANGE)
    
    generatedIp6Address = sha256GetIp6Address(timeNowMin, keyVal, constant, prefix)
    dt_object = datetime.fromtimestamp(time.time())
    print("at " + str(dt_object) + " using " + generatedIp6Address)
    
    if service == "ssh":
        command = "ssh -6 " +  generatedIp6Address
        os.system(command)
    else:
        print("This service is not supported yet. This program supports only ssh service now.")

if __name__ == "__main__":
    main()
