# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.


"""Telnet-based shell."""

# twisted imports
from twisted.protocols import telnet
from twisted.internet import protocol
from twisted.python import log, failure

# system imports
import string, copy, sys
from cStringIO import StringIO


class Shell(telnet.Telnet):
"""A Python command-line shell."""

def connectionMade(self):
telnet.Telnet.connectionMade(self)
self.lineBuffer = []

def loggedIn(self):
self.transport.write(">>> ")

def checkUserAndPass(self, username, password):
return ((self.factory.username == username) and (password == self.factory.password))

def write(self, data):
"""Write some data to the transport.
"""
self.transport.write(data)

def telnet_Command(self, cmd):
if self.lineBuffer:
if not cmd:
cmd = string.join(self.lineBuffer, '\n') + '\n\n\n'
self.doCommand(cmd)
self.lineBuffer = []
return "Command"
else:
self.lineBuffer.append(cmd)
self.transport.write("... ")
return "Command"
else:
self.doCommand(cmd)
return "Command"

def doCommand(self, cmd):

# TODO -- refactor this, Reality.author.Author, and the manhole shell
#to use common functionality (perhaps a twisted.python.code module?)
fn = '$telnet$'
result = None
try:
out = sys.stdout
sys.stdout = self
try:
code = compile(cmd,fn,'eval')
result = eval(code, self.factory.namespace)
except:
try:
code = compile(cmd, fn, 'exec')
exec code in self.factory.namespace
except SyntaxError, e:
if not self.lineBuffer and str(e)[:14] == "unexpected EOF":
self.lineBuffer.append(cmd)
self.transport.write("... ")
return
else:
failure.Failure().printTraceback(file=self)
log.deferr()
self.write('\r\n>>> ')
return
except:
io = StringIO()
failure.Failure().printTraceback(file=self)
log.deferr()
self.write('\r\n>>> ')
return
finally:
sys.stdout = out

self.factory.namespace['_'] = result
if result is not None:
self.transport.write(repr(result))
self.transport.write('\r\n')
self.transport.write(">>> ")



class ShellFactory(protocol.Factory):
username = "admin"
password = "admin"
protocol = Shell
service = None

def __init__(self):
self.namespace = {
'factory': self,
'service': None,
'_': None
}

def setService(self, service):
self.namespace['service'] = self.service = service

def __getstate__(self):
"""This returns the persistent state of this shell factory.
"""
dict = self.__dict__
ns = copy.copy(dict['namespace'])
dict['namespace'] = ns
if ns.has_key('__builtins__'):
del ns['__builtins__']
return dict