# sshserver.py - ssh protocol server support for mercurial## Copyright 2005-2007 Matt Mackall <mpm@selenic.com># Copyright 2006 Vadim Gelfer <vadim.gelfer@gmail.com>## This software may be used and distributed according to the terms of the# GNU General Public License version 2, incorporated herein by reference. from i18n import _from node import bin, heximport streamclone, util, hookimport os, sys, tempfile, urllib class sshserver(object):    def __init__(self, ui, repo):        self.ui = ui        self.repo = repo        self.lock = None        self.fin = sys.stdin        self.fout = sys.stdout         hook.redirect(True)        sys.stdout = sys.stderr         # Prevent insertion/deletion of CRs        util.set_binary(self.fin)        util.set_binary(self.fout)     def getarg(self):        argline = self.fin.readline()[:-1]        arg, l = argline.split()        val = self.fin.read(int(l))        return arg, val     def respond(self, v):        self.fout.write("%d\n" % len(v))        self.fout.write(v)        self.fout.flush()     def serve_forever(self):        try:            while self.serve_one(): pass        finally:            if self.lock is not None:                self.lock.release()        sys.exit(0)     def serve_one(self):        cmd = self.fin.readline()[:-1]        if cmd:            impl = getattr(self, 'do_' + cmd, None)            if impl: impl()            else: self.respond("")        return cmd != ''     def do_lookup(self):        arg, key = self.getarg()        assert arg == 'key'        try:            r = hex(self.repo.lookup(key))            success = 1        except Exception, inst:            r = str(inst)            success = 0        self.respond("%s %s\n" % (success, r))     def do_branchmap(self):        branchmap = self.repo.branchmap()        heads = []        for branch, nodes in branchmap.iteritems():            branchname = urllib.quote(branch)            branchnodes = [hex(node) for node in nodes]            heads.append('%s %s' % (branchname, ' '.join(branchnodes)))        self.respond('\n'.join(heads))     def do_heads(self):        h = self.repo.heads()        self.respond(" ".join(map(hex, h)) + "\n")     def do_hello(self):        '''the hello command returns a set of lines describing various        interesting things about the server, in an RFC822-like format.        Currently the only one defined is "capabilities", which        consists of a line in the form:         capabilities: space separated list of tokens        '''         caps = ['unbundle', 'lookup', 'changegroupsubset', 'branchmap']        if self.ui.configbool('server', 'uncompressed'):            caps.append('stream=%d' % self.repo.changelog.version)        self.respond("capabilities: %s\n" % (' '.join(caps),))     def do_lock(self):        '''DEPRECATED - allowing remote client to lock repo is not safe'''         self.lock = self.repo.lock()        self.respond("")     def do_unlock(self):        '''DEPRECATED'''         if self.lock:            self.lock.release()        self.lock = None        self.respond("")     def do_branches(self):        arg, nodes = self.getarg()        nodes = map(bin, nodes.split(" "))        r = []        for b in self.repo.branches(nodes):            r.append(" ".join(map(hex, b)) + "\n")        self.respond("".join(r))     def do_between(self):        arg, pairs = self.getarg()        pairs = [map(bin, p.split("-")) for p in pairs.split(" ")]        r = []        for b in self.repo.between(pairs):            r.append(" ".join(map(hex, b)) + "\n")        self.respond("".join(r))     def do_changegroup(self):        nodes = []        arg, roots = self.getarg()        nodes = map(bin, roots.split(" "))         cg = self.repo.changegroup(nodes, 'serve')        while True:            d = cg.read(4096)            if not d:                break            self.fout.write(d)         self.fout.flush()     def do_changegroupsubset(self):        argmap = dict([self.getarg(), self.getarg()])        bases = [bin(n) for n in argmap['bases'].split(' ')]        heads = [bin(n) for n in argmap['heads'].split(' ')]         cg = self.repo.changegroupsubset(bases, heads, 'serve')        while True:            d = cg.read(4096)            if not d:                break            self.fout.write(d)         self.fout.flush()     def do_addchangegroup(self):        '''DEPRECATED'''         if not self.lock:            self.respond("not locked")            return         self.respond("")        r = self.repo.addchangegroup(self.fin, 'serve', self.client_url())        self.respond(str(r))     def client_url(self):        client = os.environ.get('SSH_CLIENT', '').split(' ', 1)[0]        return 'remote:ssh:' + client     def do_unbundle(self):        their_heads = self.getarg()[1].split()         def check_heads():            heads = map(hex, self.repo.heads())            return their_heads == [hex('force')] or their_heads == heads         # fail early if possible        if not check_heads():            self.respond(_('unsynced changes'))            return         self.respond('')         # write bundle data to temporary file because it can be big        tempname = fp = None        try:            fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')            fp = os.fdopen(fd, 'wb+')             count = int(self.fin.readline())            while count:                fp.write(self.fin.read(count))                count = int(self.fin.readline())             was_locked = self.lock is not None            if not was_locked:                self.lock = self.repo.lock()            try:                if not check_heads():                    # someone else committed/pushed/unbundled while we                    # were transferring data                    self.respond(_('unsynced changes'))                    return                self.respond('')                 # push can proceed                 fp.seek(0)                r = self.repo.addchangegroup(fp, 'serve', self.client_url())                self.respond(str(r))            finally:                if not was_locked:                    self.lock.release()                    self.lock = None        finally:            if fp is not None:                fp.close()            if tempname is not None:                os.unlink(tempname)     def do_stream_out(self):        try:            for chunk in streamclone.stream_out(self.repo):                self.fout.write(chunk)            self.fout.flush()        except streamclone.StreamException, inst:            self.fout.write(str(inst))            self.fout.flush()