mikrotik-export

#!/usr/bin/env python2
#-*- coding: utf-8 -*-
from __future__ import print_function

import itertools as it, operator as op, functools as ft
from subprocess import check_output, Popen, PIPE
from contextlib import contextmanager, closing
import tempfile, stat, time
import os, sys, re, logging, glob

from twisted.internet import reactor, protocol, defer, task, error
from twisted.python import log as twisted_log
from twisted.python.filepath import FilePath
from twisted.conch.endpoints import SSHCommandClientEndpoint


@contextmanager
def safe_replacement(path, mode=None):
	if mode is None:
		try: mode = stat.S_IMODE(os.lstat(path).st_mode)
		except (OSError, IOError): pass
	kws = dict( delete=False,
		dir=os.path.dirname(path), prefix=os.path.basename(path)+'.' )
	with tempfile.NamedTemporaryFile(**kws) as tmp:
		try:
			if mode is not None: os.fchmod(tmp.fileno(), mode)
			yield tmp
			if not tmp.closed: tmp.flush()
			os.rename(tmp.name, path)
		finally:
			try: os.unlink(tmp.name)
			except (OSError, IOError): pass


class DumpProtocol(protocol.Protocol):

	def connectionMade(self):
		self.finished = defer.Deferred()
		self.dump = list()

	def dataReceived(self, data):
		self.dump.append(data)

	def connectionLost(self, reason):
		reason.trap(error.ConnectionDone)
		self.finished.callback(''.join(self.dump))


@defer.inlineCallbacks
def main(reactor, *args):
	import argparse
	parser = argparse.ArgumentParser(description='Mikrotik router backup app.')

	parser.add_argument('dst_path',
		help='Path to file to write resulting dump to.'
			' Will be rotated according to specified options, if exists.')

	parser.add_argument('-r', '--router', metavar='host',
		help='IP or hostname of router device to backup.'
			' If not specified, default gateway (as shown by "ip route get") will be used.')
	parser.add_argument('-a', '--auth-file',
		metavar='path', required=True,
		help='File with "user:password" inside. Must be specified.')
	parser.add_argument('-c', '--command',
		metavar='backup-command', default='/export',
		help='Command line to issue to get the backup (default: %(default)s).')
	parser.add_argument('-z', '--compress',
		nargs='?', default=False,
		help='Compress resulting file with gzip (default)'
			' or any other binary specified with this option.')
	parser.add_argument('-n', '--rotate', type=int, default=9,
		help='Number of old files to keep around with date-suffix.'
			'Setting to 0 will disable the feature, i.e. only latest file will be available.')

	parser.add_argument('--debug', action='store_true', help='Verbose operation mode.')
	opts = parser.parse_args(args)

	if opts.compress is None: opts.compress = 'gzip'
	compress = opts.compress

	global log
	logging.basicConfig(level=logging.DEBUG if opts.debug else logging.WARNING)
	twisted_log.PythonLoggingObserver(loggerName='twisted').start()
	log = logging.getLogger()

	with open(opts.auth_file, 'rb') as src:
		user, pw = src.read().strip().split(':', 1)

	host = opts.router
	if not host:
		route_dst = '1.2.3.4'
		route = check_output(['ip', 'ro', 'get', route_dst])
		for line in route.splitlines():
			match = re.search(( r'^\s*{} via (?P<gw>[\d.]+).*'
				' dev (?P<dev>\S+).* src (?P<src>[\d.]+)' ).format(route_dst), line)
			if match: break
		else:
			parser.error('Unable to determine default router (via "ip route get")')
		host, dev, src = (match.group(k) for k in ['gw', 'dev', 'src'])

	log.debug('Connecting to %s@%s', user, host)
	ep = SSHCommandClientEndpoint\
		.newConnection(reactor, opts.command, user, host, password=pw)
	proto = yield ep.connect(protocol.Factory.forProtocol(DumpProtocol))
	dump = yield proto.finished
	log.debug('Got dump data (size: %sB)', len(dump))

	with safe_replacement(opts.dst_path) as tmp:
		if compress:
			log.debug('Using compression binary: %s', compress)
			compress = Popen([compress], stdin=PIPE, stdout=tmp)
			tmp = compress.stdin
		tmp.write(dump)
		if compress:
			compress.stdin.close()
			err = compress.wait()
			if err != 0:
				log.error( 'Compression command (%s)'
					' exited with non-zero code: %s', opts.compress, err )
				raise SystemExit(1)
		if opts.rotate > 0:
			p = FilePath(opts.dst_path)
			if p.exists():
				ts = p.getModificationTime()
				p.moveTo(FilePath('{}.{}'.format( opts.dst_path,
					time.strftime('%Y-%m-%d_%H%M%S', time.localtime(ts)) )))
			ps = map(FilePath, sorted(glob.glob('{}.*'.format(opts.dst_path)), reverse=True))
			while len(ps) > (opts.rotate + 1): ps.pop().remove() # +1 for tmp file

if __name__ == '__main__': task.react(main, sys.argv[1:])