import binascii
from scapy.all import *
import socket

memcard = {
  0 : "No Memory Card",
  1 : "SPRAM",
  2 : "EPROM",
  3 : "EEPROM"
}

def memory_card(value):
  mem_card = memcard[value] if value in memcard else "Unknown Memory Card Type"
  return mem_card


def send_tcp(s):
	req_addr = binascii.unhexlify("46494e530000000c000000000000000000000000")
	controller_data_read = binascii.unhexlify("46494e5300000015000000020000000080000200")
	controller_data_read2 = binascii.unhexlify("000000ef050501")
	BUFFER_SIZE = 1024
	try:
		s.send(req_addr)
		response = s.recv(BUFFER_SIZE)
	except:
		return False
	header = response[0]
	if(header == 70):
		address = response[:23]
		controller_data = controller_data_read+address+controller_data_read2+binascii.unhexlify('00')
		try:
			s.send(controller_data)
			res = s.recv(BUFFER_SIZE)
		except:
			return False
		return res
	return "ERROR"

def send_udp(host,port):
	s = socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
	BUFFER_SIZE = 1024
	controller_data_read = binascii.unhexlify("800002000000006300ef050100")
	try:
		s.sendto(controller_data_read,(host,port))
		response = s.recv(BUFFER_SIZE)
	except:
		return False
	return response


def action(host,port,protocol):
	output={}
	s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	try:
		s.settimeout(3)
		s.connect((host,port))
	except:
		return False
	response = ""
	offset = 0
	if (protocol == "tcp"):
		offset = 16
		response = send_tcp(s)
	else:
		response = send_udp(host,port)
	header = response[0]
	if(header==192 or header==193 or header==70):
		response_code = response[12+offset]
		if(response_code == 2081):
			output["Response Code"] = "Data cannot be changed (0x2108)"
		elif(response_code == 290):
			output["Response Code"] = "The mode is wrong (executing) (0x2201)"
		elif(response_code == 0):
			output["Response Code"] = "Normal completion (0x0000)"
			output["Controller Model"] = ((response[14+offset:].split(b"\x00",1))[0]).decode()
			output["Controller Version"] = ((response[34+offset:].split(b"\x00",1))[0]).decode()
			output["For System Use"] = ((response[54+offset:].split(b"\x00",1))[0]).decode()
			output["Program Area Size"] = str(response[95])
			output["IOM size"] = str(response[96])
			output["No. DM Words"] = str(int(binascii.hexlify(response[97:99]),base=16))
			output["Timer/Counter"] = str(response[99])
			output["Expansion DM Size"] = str(response[100])
			output["No. of steps/transitions"] = str(int(binascii.hexlify(response[101:102]),base=16))
			mem_card_type = response[102]
			output["Kind of Memory Card"] = memory_card(mem_card_type)
			output["Memory Card Size"] = str(int(binascii.hexlify(response[103:104]),base=16))
			return (output)

def get_info(ip,port,protocol='udp'):
	return(action(ip,port,protocol))