import socket
import threading
import sys
import os
import base64
import random
import ssl
import traceback
import core
import pexpect
import re
from colorama import *
from Crypto.Cipher import AES
from Crypto import Random
from datetime import datetime
GOOD = Fore.GREEN + " + " + Fore.RESET
BAD = Fore.RED + " - " + Fore.RESET
WARN = Fore.YELLOW + " * " + Fore.RESET
INFO = Fore.BLUE + " + " + Fore.RESET
prompt = Fore.BLUE + ">> " + Fore.RESET
SERVER_PORT = 9999
SERVER_IP = "127.0.0.1"
ERR_MSG = BAD + "Lost server connection. Please try again later."
SEP = "|:|"
END_SEP = "~"
####################################################################
#
# Message Types
#
####################################################################
MSG = 0
GET = 1
CLOSE = 2
INFO = 3
REQ = 4
CLI_INIT = 10
CLI = 11
CLI_RESP = 12
BS = 16
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS)
unpad = lambda s : s[0:-ord(s[-1])]
replace_seps = lambda m: m.replace("|::|", SEP).replace("!::!", END_SEP)
save_seps = lambda m: m.replace(SEP, "|::|").replace(END_SEP, "!::!")
def get_date():
return "%02d:%02d:%02d" % (datetime.now().hour, datetime.now().minute, datetime.now().second)
class PinaColadaSocket(object):
def __init__(self, name, target_port, server_ip):
self.port = target_port
self.ip = server_ip
self.name = name
self.socket = None
self.core = core.PinaColada()
self.keys = {}
self.conn = None
self.cli = None
print "[*] Attempting to connect to server"
def connect(self):
client = None
try:
#s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
#s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
#client = ssl.wrap_socket(s, ca_certs="cert/server.crt", cert_reqs=ssl.CERT_REQUIRED)
client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
client.connect((self.ip, self.port))
self.socket = client
shared_prime, shared_base = client.recv(10).split("|")
shared_prime = int(shared_prime)
shared_base = int(shared_base.replace("~", ""))
client_secret = random.randint(0, 99)
a = long(client.recv(1024).replace("~", ""))
b = (shared_base**client_secret) % shared_prime
client.send("%ld" % b)
self.keys[client] = pad("%ld" % ((a ** client_secret) % shared_prime))
client.settimeout(None) # Remove the timeout
self.send(MSG, "0", "PinaColada")
print("[*] Successfully connected.")
self.receive(client)
except Exception as e:
print e
print traceback.print_exc()
raw_input("")
finally:
self.shutdown()
def shutdown(self):
print GOOD + "Exiting..."
####################################################################
#
# Message Handling
#
####################################################################
def get(self, request):
reqs = request.split()
if request == "help":
return "help message" # TODO
if request == "categories":
return self.core.get_categories()
if "capabilities" in request:
return self.core.get_capabilities(category=reqs[1]) # TODO POTENTIAL BUG
if request == "ip":
return self.core.get_local_ip(self.core.default_iface) # TODO ADD INTERFACES
if request == "list":
info = {}
for cat in self.core.get_categories():
caps = self.core.get_capabilities(cat)
if caps != []:
info[cat] = []
for cap in caps:
info[cat].append(cap)
return info
def cli_init(self):
if self.cli:
self.cli.terminate(force=True)
self.cli = pexpect.spawn("sudo python cli.py")
self.cli.expect(re.escape(prompt))
return self.cli.before + prompt
def cli_communicate(self, data):
self.cli.sendline(data)
self.cli.expect(re.escape(prompt))
return "\n".join(self.cli.before.split("\n")[1:]) + prompt
def send(self, message_type, requester, data):
#print "SENDING: <%d, %s>" %(message_type, data)
self.socket.send(self.encrypt(self.pack_data(message_type, requester, data), self.socket) + "\n")
def encrypt(self, string, sock):
iv = Random.new().read(AES.block_size)
cipher = AES.new(self.keys[sock], AES.MODE_CBC, iv)
return base64.b64encode(iv + cipher.encrypt(pad(string)))
def decrypt(self, msg, sock):
enc = base64.b64decode(msg)
cipher = AES.new(self.keys[sock], AES.MODE_CBC, enc[:16])
return unpad(cipher.decrypt(enc[16:]))
def unpack_data(self, msg):
msgs = [replace_seps(s) for s in msg.split(SEP)]
try:
msgs[0] = int(msgs[0]) # convert type to int
finally:
return msgs
def pack_data(self, message_type, name, data):
return str(message_type) + SEP + save_seps(name) + SEP + save_seps(data) + END_SEP
def print_msg(self, message):
print message
def handle(self, data):
message_type, name, d = self.unpack_data(data)
print "DATA FOR CLI: %s" % d
message_type = int(message_type)
print message_type, name, d
if message_type == CLI_INIT:
self.send(CLI_RESP, name, self.cli_init())
if message_type == CLI:
self.send(CLI_RESP, name, self.cli_communicate(d))
self.print_msg(data)
####################################################################
#
# Receive Handling
#
####################################################################
def receive(self, sock):
try:
while True:
data = sock.recv(1024)
msgs = filter(None, self.decrypt(data, sock).split(END_SEP))
for m in msgs:
self.handle(m)
except Exception as e:
print e
traceback.print_exc()
sock.close()
finally:
sock.close()
self.shutdown()
if __name__ == "__main__":
if os.getuid() != 0:
print BAD + "Please run me as root!"
sys.exit()
PinaColadaSocket("Client", SERVER_PORT, SERVER_IP).connect()
