Add DCC module

This commit is contained in:
dave 2018-01-16 17:04:15 -08:00
parent cc8951f865
commit d959db127f
8 changed files with 278 additions and 3 deletions

41
docs/api/modules/dcc.rst Normal file
View File

@ -0,0 +1,41 @@
:mod:`DCC` --- Interface to DCC file transfers
==============================================
A module providing a high-level interface for DCC file transfers.
DCC file transfers involve having the sender listen on some tcp port and the receiver connect to the port to initiate
the file transfer. The file name and length as well as the tcp port and address are shared via other means.
Config
------
.. code-block:: json
{
"port_range": [40690, 40990],
"public_addr": "127.0.0.1",
"bind_host": "127.0.0.1"
}
.. cmdoption:: port_range
The range of ports between which arbitrary ports will be used for file transfers
.. cmdoption:: public_addr
When sending files, what address we will advertise as being connectable on
.. cmdoption:: bind_host
What IP address to bind to when creating listener sockets for the file send role.
Class Reference
---------------
.. automodule:: pyircbot.modules.DCC
:members:
:undoc-members:
:show-inheritance:

2
examples/.gitignore vendored Normal file
View File

@ -0,0 +1,2 @@
# instance config files
/config_*.json

169
pyircbot/modules/DCC.py Normal file
View File

@ -0,0 +1,169 @@
"""
.. module:: DCC
:synopsis: Module providing support for IRC's dcc protocol
.. moduleauthor:: Dave Pedu <dave@davepedu.com>
"""
import os
from pyircbot.modulebase import ModuleBase
import socket
from threading import Thread
from random import randint
from time import sleep
BUFFSIZE = 8192
def ip2int(ipstr):
"""
Convert an ip address string to an integer
"""
num = 0
for octet in ipstr.split("."):
num = num << 8 | int(octet)
return num
def int2ip(num):
"""
Convert an integer to an ip address string
"""
octs = []
for octet in range(0, 4):
octs.append(str((num & (255 << (8 * octet))) >> (8 * octet)))
return ".".join(octs[::-1])
class DCC(ModuleBase):
def __init__(self, bot, name):
super().__init__(bot, name)
self.services = ["dcc"]
self.is_kill = False
self.transfers = []
def offer(self, file_path, port=None):
"""
Offer a file to another user.
- check file size
- start listener socket thread on some port
- info about the file: tuple of (ip, port, size)
"""
port_range = self.config.get("port_range", [40000, 60000]) # TODO it would be better to let the system assign
port = randint(*port_range)
bind_addr = self.config.get("bind_host", "0.0.0.0")
advertise_addr = self.config.get("public_addr", bind_addr)
flen = os.path.getsize(file_path)
offer = OfferThread(self, file_path, bind_addr, port) # offers are considered ephemeral. even if this module is
# unloaded, initiated transfers may continue. They will not block python from exiting (at which time they *will*
# be terminated).
offer.start()
return (ip2int(advertise_addr), port, flen, offer)
def recieve(self, host, port, length):
"""
Receive a file another user has offered. Returns a generator that yields data chunks.
"""
return RecieveGenerator(host, port, length)
class RecieveGenerator(object):
def __init__(self, host, port, length):
self.host = host
self.port = port
self.length = length
def __iter__(self):
self.sock = socket.create_connection((self.host, self.port), timeout=10)
total = 0
try:
while True:
if total == self.length:
break
chunk = self.sock.recv(BUFFSIZE)
total += len(chunk)
if not chunk:
break
yield chunk
if total >= self.length:
break
if total != self.length:
raise Exception("Transfer failed: expected {} bytes but got {}".format(self.length, total))
raise StopIteration()
finally:
self.sock.shutdown(socket.SHUT_RDWR)
self.sock.close()
class OfferThread(Thread):
def __init__(self, master, path, bind_addr, port, timeout=30):
"""
DCC file transfer offer listener
:param master: reference to the parent module
:param path: file path to be opened and transferred
:param bind_addr: address str to bind the listener socket to
:param port: port number int to listen on
:param timeout: number of seconds to give up after
"""
super().__init__()
self.master = master
self.path = path
self.bind = bind_addr
self.port = port
self.timeout = timeout
self.listener = None
self.daemon = True
self.bound = False
Thread(target=self.abort, daemon=True).start()
def run(self):
"""
Open a server socket that accepts a single connections. When the first client connects, send the contents of the
offered file.
"""
self.listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
self.listener.bind((self.bind, self.port))
self.listener.listen(1)
self.bound = True
(clientsocket, address) = self.listener.accept()
try:
self.send_file(clientsocket)
finally:
clientsocket.shutdown(socket.SHUT_RDWR)
clientsocket.close()
finally:
self.listener.shutdown(socket.SHUT_RDWR)
self.listener.close()
def abort(self):
"""
Expire the offer after a timeout.
"""
sleep(self.timeout)
self.stopoffer()
def send_file(self, socket):
"""
Send the contents of the offered file to the passed socket
:param socket: socket object ready for sending
:type socket: socket.socket
"""
with open(self.path, 'rb') as f:
while not self.master.is_kill:
chunk = f.read(BUFFSIZE)
if not chunk:
break
socket.send(chunk)
def stopoffer(self):
"""
Prematurely shut down & cleanup the offer socket
"""
try:
self.listener.shutdown(socket.SHUT_RDWR)
self.listener.close()
except Exception: # should only error if already cleaned up
pass

View File

@ -15,7 +15,7 @@ jaraco.classes==1.4.3
jedi==0.10.2
lxml==4.1.1
mock==2.0.0
msgbus==0.0.1
-e git+http://gitlab.davepedu.com/dave/pymsgbus.git#egg=msgbus
packaging==16.8
pbr==3.1.1
pexpect==4.2.1

View File

@ -1,8 +1,9 @@
#!/bin/bash
CONFPATH=${1:-examples/config.json}
shift || true
export PYTHONUNBUFFERED=1
export PYTHONPATH=.
./bin/pyircbot -c $CONFPATH --debug
./bin/pyircbot -c $CONFPATH --debug $@

View File

@ -3,4 +3,4 @@
export PYTHONUNBUFFERED=1
export PYTHONPATH=.
py.test --cov=pyircbot --cov-report html -n 4 tests/
py.test --fulltrace --cov=pyircbot --cov-report html -n 4 tests/

51
tests/modules/test_dcc.py Normal file
View File

@ -0,0 +1,51 @@
import pytest
import os
import hashlib
from time import sleep
from pyircbot.modules.DCC import int2ip
from tests.lib import * # NOQA - fixtures
@pytest.fixture
def dccbot(fakebot):
"""
Provide a bot loaded with the DCC module
"""
fakebot.botconfig["module_configs"]["DCC"] = {
"port_range": [40690, 40990],
"public_addr": "127.0.0.1",
"bind_host": "127.0.0.1"
}
fakebot.loadmodule("DCC")
return fakebot
def test_offerrecv(dccbot, tmpdir):
# allocate a temp file
flen = 1024 * 51
tmpfpath = os.path.join(tmpdir.dirname, "hello.bin")
with open(tmpfpath, 'wb') as fout:
fout.write(os.urandom(flen))
# hash the tmpfile for later comparison
m = hashlib.sha256()
with open(tmpfpath, "rb") as ftmp:
m.update(ftmp.read())
srchash = m.hexdigest()
# offer th file over DCC
ip, port, reported_len, offer = dccbot.moduleInstances['DCC'].offer(tmpfpath)
reported_len = int(reported_len)
assert reported_len == flen, "offer reported wrong file length!"
ip = int2ip(ip)
while not offer.bound:
sleep(0.001)
# receive the file over DCC
print("connecting to {}:{}".format(ip, port))
recver = dccbot.moduleInstances['DCC'].recieve(ip, port, reported_len)
data = b''
d = hashlib.sha256()
for chunk in iter(recver):
data += chunk
d.update(chunk)
# verify hashes and lengths
assert len(data) == flen, "file not completely transferred"
assert d.hexdigest() == srchash, "file was mangled in transfer"

11
tests/test_common.py Normal file
View File

@ -0,0 +1,11 @@
from pyircbot import common
def test_parse_line():
assert common.parse_irc_line(":chuck!~chuck@foobar PRIVMSG #jesusandhacking :asdf") == \
('PRIVMSG', ['#jesusandhacking'], 'chuck!~chuck@foobar', "asdf")
def test_parse_notrailing():
assert common.parse_irc_line(":chuck!~chuck@foobar MODE #jesusandhacking -o asciibot") == \
('MODE', ['#jesusandhacking', '-o', 'asciibot'], 'chuck!~chuck@foobar', None)