import socket
import sys
import argparse

message = "test"


def main(argv):
    parser = argparse.ArgumentParser(
        description='Send socket message to ip and port with protocol.')
    parser.add_argument('-i', '--ip',
                        default='localhost',
                        dest='ip',
                        help='Provide destination ip. Defaults to localhost',
                        type=str
                        )
    parser.add_argument('-p', '--port',
                        default=80,
                        dest='port',
                        help='Provide destination port. Defaults to 80',
                        type=int
                        )
    parser.add_argument('-t', '--type', choices=['UDP', 'TCP'],
                        default='UDP',
                        dest='proto',
                        help='Provide protocol. Defaults to UDP',
                        type=str
                        )
    args = parser.parse_args()

    print(f'Sending message to {args.ip}:{args.port} width {args.proto}')

    if args.proto == 'TCP':
        sock = socket.socket()
        try:
            sock.connect((args.ip, args.port))
            sock.send(message.encode())
            print("Message sent.")
            sock.close()
        except:
            print("Cannot send message...")

    else:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            sock.sendto(message.encode('utf-8'), (args.ip, args.port))
            print("Message sent.")
            sock.close()
        except:
            print("Cannot send message...")


if __name__ == "__main__":
    main(sys.argv[1:])