# Trixy: Create network listeners, tunnels, and outbound connections in a
# modular way allowing interception and modification of the traffic.
#
# Copyright (C) 2014-2015 Austin Hartzheim
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# 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, see <http://www.gnu.org/licenses/>.
import asyncio
import socket
[docs]class TrixyNode():
'''
A base class for TrixyNodes that implements some default packet
forwarding and node linking.
'''
def __init__(self):
self.downstream_nodes = []
self.upstream_nodes = []
[docs] def add_downstream_node(self, node):
'''
Add a one direction downstream link to the node parameter.
:param TrixyNode node: The downstream node to create a
unidirectional link to.
'''
if node not in self.downstream_nodes:
self.downstream_nodes.append(node)
[docs] def add_upstream_node(self, node):
'''
Add a one direction upstream link to the node parameter.
:param TrixyNode node: The upstream node to create a
unidirectional link to.
'''
if node not in self.upstream_nodes:
self.upstream_nodes.append(node)
[docs] def connect_node(self, node):
'''
Create a bidirectional connection between the two nodes with
the downstream node being the parameter.
:param TrixyNode node: The downstream node to create a
bidirectional connection to.
'''
self.add_upstream_node(node)
node.add_downstream_node(self)
[docs] def forward_packet_down(self, data):
'''
Forward data to all downstream nodes.
:param bytes data: The data to forward.
'''
for node in self.downstream_nodes:
node.handle_packet_down(data)
[docs] def forward_packet_up(self, data):
'''
Forward data to all upstream nodes.
:param bytes data: The data to forward.
'''
for node in self.upstream_nodes:
node.handle_packet_up(data)
[docs] def handle_close(self, direction='down'):
'''
The connection has closed on one end. So, shutdown what we are
doing and notify the nodes we are connected to.
:param str direction: 'down' or 'up' depending on if downstream
nodes need to be closed, or upstream nodes need to be closed.
'''
if direction == 'down':
for node in self.downstream_nodes:
node.handle_close(direction='down')
elif direction == 'up':
for node in self.upstream_nodes:
node.handle_close(direction='up')
[docs] def handle_packet_down(self, data):
'''
Hadle data moving downwards. TrixyProcessor children should
perform some action on `data` whereas `TrixyOutput` children
should send the data to the desired output location.
Generally, the a child implementation of this method should
be implemented such that it calls self.forward_packet_down
with the data (post-modification if necessary) to forward the
data to other processors in the chain. However, if the
processor is a filter, it may drop the packet by omitting that
call.
:param bytes data: The data that is being handled.
'''
self.forward_packet_down(data)
[docs] def handle_packet_up(self, data):
'''
Hadle data moving upwards. TrixyProcessor children should
perform some action on `data` whereas `TrixyOutput` children
should send the data to the desired output location.
Generally, the a child implementation of this method should
be implemented such that it calls self.forward_packet_down
with the data (post-modification if necessary) to forward the
data to other processors in the chain. However, if the
processor is a filter, it may drop the packet by omitting that
call.
:param bytes data: The data that is being handled.
'''
self.forward_packet_up(data)
[docs]class TrixyServer():
'''
Main server to grab incoming connections and forward them.
'''
# TODO: evaluate if there is a place for a server class or if
# asyncio should be used (loop.create_server) instead.
def __init__(self, tinput, host, port, loop=None):
'''
:param TrixyInput tinput: instantiated every time an incoming
connection is grabbed.
:param str host: the hostname to bind to.
:param int port: the port number to bind to.
:param loop: you may optionally specify your own event loop.
'''
super().__init__()
if not loop:
loop = asyncio.get_event_loop()
self.loop = loop
self.tinput = tinput
self.host = host
self.port = port
self.managing_loop = False
self.setup_socket()
def setup_socket(self):
coro = self.loop.create_server(lambda: self.tinput(self.loop),
self.host, self.port)
self.server = self.loop.run_until_complete(coro)
[docs] def close(self):
'''
Shutdown the server.
'''
self.server.close()
self.loop.run_until_complete(self.server.wait_closed())
[docs] def run_loop(self):
'''
Run the event loop so that the server functions.
'''
self.managing_loop = True
self.loop.run_forever()
def connection_lost(self):
self.server.close()
self.loop.run_until_complete(self.server.wait_closed())
if self.managing_loop:
self.loop.close()
[docs]class TrixyProcessor(TrixyNode):
'''
Perform processing on data moving through Trixy.
'''
pass
[docs]class TrixyOutput(TrixyNode, asyncio.Protocol):
'''
Output node; generally to connect to another network service.
'''
def __init__(self, loop, host, port, autoconnect=True):
'''
:param loop: The asyncio event loop.
'''
super().__init__()
self.loop = loop
self.host = host
self.port = port
self.transport = None
self.buffer = b''
if autoconnect:
self.connect()
[docs] def connect(self):
'''
Connect to the saved hsot and port. This method is commonly
called by __init__() when `autoconnect` is enabled.
'''
coro = self.loop.create_connection(lambda: self, self.host, self.port)
self.task = asyncio.async(coro)
[docs] def forward_packet_up(self, data):
'''
If the transport is connected, send the data onwards; otherwise
it will be added to the buffer.
This method require buffering because data often becomes
available immediately for upload, but the transport needs time
to connect. (This is especially the case when autoconnect is
set to True).
'''
if self.transport is not None:
self.transport.write(data)
else:
self.buffer += data
[docs] def data_received(self, data):
'''
Received incoming data; forward it down the chain.
:param bytes data: a bytes object of the received data.
'''
self.forward_packet_down(data)
[docs] def connection_made(self, transport):
'''
Outbound connection successful; save the transport.
'''
self.transport = transport
self.transport.write(self.buffer)