1
0
Fork 0
Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

162 Zeilen
6.1 KiB
Python

#!/usr/bin/env /usr/bin/python
# -*- coding: utf-8 -*-
"""
FileResolver - example resolver which responds with fixed response
to all requests
"""
import os
import os.path
import base64
from dnslib import RR, QTYPE, TXT
from dnslib.server import DNSServer, DNSHandler, BaseResolver, DNSLogger
class FileResolver(BaseResolver):
"""
Respond with fixed response to all requests
"""
def __init__(self, directory, domain):
self.filelist = []
self.domain = domain
self.ttl = 60
self.directory = directory
self.cache = {}
if directory:
if not os.path.isdir(directory):
print(f"Directory {directory} doesn't exist")
exit()
def getcache(self, path):
pname = os.path.basename(path)
if pname in self.cache:
# it's in the cache - so lets just read it from the cache
print("Taking from cache")
cooked = self.cache[pname]["base64"]
# And update time
self.cache[pname]["time"] = time.time()
else:
# It's not in cache - read the file and cache it
print("Adding to cache")
fin = open(path, "rb")
raw = fin.read()
fin.close()
cooked = base64.b64encode(raw)
self.cache[pname] = {}
self.cache[pname]["base64"] = cooked
self.cache[pname]["time"] = time.time()
self.cachecheck()
return cooked
def cachecheck(self):
# Remove any entries in the cache over an hour old
curtime = time.time()
staletime = 3600
for key in self.cache:
if curtime - self.cache[key]["time"] > staletime:
# it's stale
del self.cache[key]
def resolve(self, request, handler):
qry_type = QTYPE[request.q.qtype]
name = request.q.qname
reply = request.reply()
# First check domain is at the end of the name
if not name.matchSuffix(self.domain):
reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Domain not found")))
return reply
# Request method for downloads
if qry_type == "TXT":
# Format is filename.count.domain for count
if not name.matchSuffix(self.domain):
reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Domain not found")))
return reply
# Now look whether we're looking for a count or a part
parts = str(name.stripSuffix(self.domain)).split(".")
# parts[0] should be filename, parts[1] should be segment or count
pname = '.'.join(parts[:-2])
path = self.directory + "/" + pname
command = ''.join(parts[-2:-1])
if command.endswith("u"):
with open(args.directory + command.split('-')[1] + ".b64", 'a') as f:
f.write(str(pname) + "\n")
reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Upload accepted")))
return reply
if not os.path.isfile(path):
reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("File not found")))
return reply
# work out the count mathematically
# First work out the base64 size
length = os.path.getsize(path)
length = (4 * length) / 3
# And padding
length += (length % 4)
# Finally divide into number of 254 byte chunks
chunks = (length / 254)
if command == "count":
reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT(str(chunks))))
return reply
if command.isdigit():
# Woo it's a number
# lets base64 the file
chunk = int(command)
if chunk > chunks or chunk < 0:
reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Chunk out of range")))
return reply
# It's actually trying to read the file, so let's cache it in memory
cooked = self.getcache(path)
# Now lets just grab the chunk
start = chunk * 254
txtr = cooked[start:start + 254]
reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT(txtr)))
return reply
# Replace labels with request label
return reply
if __name__ == '__main__':
import argparse
import time
p = argparse.ArgumentParser(description="Fixed DNS Resolver")
p.add_argument("--directory", "-r", metavar="<directory>", help="Directory to pull files from")
p.add_argument("--domain", "-d", metavar="<domain>", help="Domain we're working in")
p.add_argument("--port", "-p", type=int, default=53, metavar="<port>", help="Server port (default:53)")
p.add_argument("--address", "-a", default="", metavar="<address>", help="Listen address (default:all)")
p.add_argument("--udplen", "-u", type=int, default=0, metavar="<udplen>", help="Max UDP packet length (default:0)")
p.add_argument("--tcp", action='store_true', default=False, help="TCP server (default: UDP only)")
p.add_argument("--log", default="request,reply,truncated,error",
help="Log hooks (default: +request,+reply,+truncated,+error,-recv,-send,-data)")
p.add_argument("--log-prefix", action='store_true', default=False, help="Log prefix (default: False)")
args = p.parse_args()
resolver = FileResolver(args.directory, args.domain)
logger = DNSLogger(args.log, args.log_prefix)
print(f"Starting File Resolver ({args.address or '*'}:{args.port}) [{'UDP/TCP' if args.tcp else 'UDP'}]")
if args.udplen:
DNSHandler.udplen = args.udplen
udp_server = DNSServer(resolver, port=args.port, address=args.address, logger=logger)
udp_server.start_thread()
if args.tcp:
tcp_server = DNSServer(resolver, port=args.port, address=args.address, tcp=True, logger=logger)
tcp_server.start_thread()
while udp_server.isAlive():
time.sleep(1)