Merge pull request #21 from sorend/feature/add_before_hosts_arg

Feature/add before hosts arg

README.md

@@ -43,6 +43,8 @@ c.add("newsvu", Hostname="ssh-new.svu.local", Port=22, User="stud1234",
                 RemoteForward=["localhost:2022 localhost:22", "localhost:2025 localhost:25"])
 print("newsvu", c.host("newsvu"))
 
+c.add("oldsvu", before_host="newsvu", Hostname="ssh-old.svu.local", Port=22, User="Stud1234")
+
 c.rename("newsvu", "svu-new")
 print("svu-new", c.host("svu-new"))
 

sshconf.py

@@ -155,14 +155,22 @@ def _indent(s):
     return s[0: len(s) - len(s.lstrip())]
 
 
+def _find_insert_idx(before_host, lines):
+    first_host_idx = next(idx for idx, x in enumerate(lines) if x.host == before_host)
+    for i in reversed(range(first_host_idx)):
+        if lines[i].host is not None:
+            return i+1
+    return 0
+
+
 class SshConfigFile(object):
     """
     Class for manipulating SSH configuration.
     """
 
     def __init__(self, lines):
         self.lines_ = []
-        self.hosts_ = set()
+        self.hosts_ = []
         self.parse(lines)
 
     def parse(self, lines):
@@ -175,7 +183,7 @@ def parse(self, lines):
                 key, value = kv_
                 if key.lower() == "host":
                     cur_entry = value
-                    self.hosts_.add(value)
+                    self.hosts_.append(value)
                 else:
                     indents.append(_indent(line))
                 self.lines_.append(ConfigLine(line=line, host=cur_entry, key=key, value=value))
@@ -300,10 +308,9 @@ def rename(self, old_host, new_host):
                 if line.key.lower() == "host":
                     line.value = new_host
                     line.line = "Host %s" % new_host
-        self.hosts_.remove(old_host)  # update host cache
-        self.hosts_.add(new_host)
+        self.hosts_[self.hosts_.index(old_host)] = new_host  # update host cache
 
-    def add(self, host, **kwargs):
+    def add(self, host, before_host=None, **kwargs):
         """
         Add another host to the SSH configuration.
 
@@ -314,17 +321,36 @@ def add(self, host, **kwargs):
         """
         if host in self.hosts_:
             raise ValueError("Host %s: exists (use update)." % host)
-        self.hosts_.add(host)
-        self.lines_.append(ConfigLine(line="", host=None))
-        self.lines_.append(ConfigLine(line="Host %s" % host, host=host, key="Host", value=host))
-        for k, v in kwargs.items():
-            if type(v) not in [list, tuple]:
-                v = [v]
-            mapped_k = _remap_key(k)
-            for value in v:
-                new_line = self._new_line(mapped_k, value)
-                self.lines_.append(ConfigLine(line=new_line, host=host, key=mapped_k, value=value))
-        self.lines_.append(ConfigLine(line="", host=None))
+
+        if before_host is not None and before_host not in self.hosts_:
+            raise ValueError("Host %s: not found, cannot insert before it." % host)
+
+        def kwargs_to_lines(kv):
+            lines = []
+            for k, v in kv.items():
+                if type(v) not in [list, tuple]:
+                    v = [v]
+                mapped_k = _remap_key(k)
+                for value in v:
+                    new_line = self._new_line(mapped_k, value)
+                    lines.append(ConfigLine(line=new_line, host=host, key=mapped_k, value=value))
+            return lines
+
+        new_lines = [
+            ConfigLine(line="", host=None),
+            ConfigLine(line="Host %s" % host, host=host, key="Host", value=host)
+        ] + kwargs_to_lines(kwargs) + [
+            ConfigLine(line="", host=None)
+        ]
+
+        if before_host is not None:
+            insert_idx = _find_insert_idx(before_host, self.lines_)
+            for idx, l in enumerate(new_lines):
+                self.lines_.insert(insert_idx + idx, l)
+            self.hosts_.insert(self.hosts_.index(before_host), host)
+        else:
+            self.lines_.extend(new_lines)
+            self.hosts_.append(host)
 
     def remove(self, host):
         """
@@ -336,12 +362,12 @@ def remove(self, host):
         """
         if host not in self.hosts_:
             raise ValueError("Host %s: not found." % host)
-        self.hosts_.remove(host)
         # remove lines, including comments inside the host lines
         host_lines = [idx for idx, x in enumerate(self.lines_) if x.host == host]
         remove_range = reversed(range(min(host_lines), max(host_lines) + 1))
         for idx in remove_range:
             del self.lines_[idx]
+        self.hosts_.remove(host)
 
     def config(self, filter_includes=False):
         """
@@ -407,9 +433,9 @@ def hosts(self):
         -------
         Tuple of Host entries (including "*" if found)
         """
-        hosts = set()
+        hosts = []
         for p, c in self.configs_:
-            hosts.update(c.hosts())
+            hosts.extend(c.hosts())
         return tuple(hosts)
 
     def host(self, host):
@@ -477,7 +503,7 @@ def rename(self, old_host, new_host):
             if old_host in c.hosts_:
                 c.rename(old_host, new_host)
 
-    def add(self, host, **kwargs):
+    def add(self, host, before_host=None, **kwargs):
         """
         Add another host to the SSH configuration.
 
@@ -486,7 +512,7 @@ def add(self, host, **kwargs):
         host: The Host entry to add.
         **kwargs: The parameters for the host (without "Host" parameter itself)
         """
-        self.configs_[0][1].add(host, **kwargs)
+        self.configs_[0][1].add(host, before_host=before_host, **kwargs)
 
     def remove(self, host):
         """

tests/test_sshconf.py

@@ -43,8 +43,14 @@ def test_rename():
 
     assert c.host("svu")["hostname"] == "www.svuniversity.ac.in"
 
+    hosts = c.hosts()
+
     c.rename("svu", "svu-new")
 
+    hosts2 = c.hosts()
+
+    assert hosts.index("svu") == hosts2.index("svu-new")  # order is the same
+
     assert "Host svu-new" in c.config()
     assert "Host svu\n" not in c.config()
     assert "svu" not in c.hosts()
@@ -69,9 +75,15 @@ def test_add():
 
     c = sshconf.read_ssh_config(test_config)
 
+    hosts = list(c.hosts())
+
     c.add("venkateswara", Hostname="venkateswara.onion", User="other", Port=22,
           ProxyCommand="nc -w 300 -x localhost:9050 %h %p")
 
+    hosts2 = list(c.hosts())
+
+    assert hosts + ["venkateswara"] == hosts2
+
     assert "venkateswara" in c.hosts()
     assert c.host("venkateswara")["proxycommand"] == "nc -w 300 -x localhost:9050 %h %p"
 
@@ -84,6 +96,42 @@ def test_add():
         c.add("venkateswara")
 
 
+def test_add_before_host():
+
+    c = sshconf.read_ssh_config(test_config)
+
+    hosts = list(c.hosts())
+
+    c.add("venkateswara", before_host="svu", Hostname="venkateswara.onion", User="other", Port=22,
+          ProxyCommand="nc -w 300 -x localhost:9050 %h %p")
+
+    hosts2 = list(c.hosts())
+
+    c.add("venkateswara2", before_host="*", Hostname="venkateswara.onion", User="other", Port=22,
+          ProxyCommand="nc -w 300 -x localhost:9050 %h %p")
+
+    hosts3 = list(c.hosts())
+
+    new_config = c.config()
+
+    print("new config", new_config)
+
+    assert hosts == ["*", "svu"]
+    assert hosts2 == ["*", "venkateswara", "svu"]
+    assert hosts3 == ["venkateswara2", "*", "venkateswara", "svu"]
+
+    assert "venkateswara" in c.hosts()
+
+    assert "Host venkateswara" in new_config
+
+    with pytest.raises(ValueError):  # cant add before a host that is not found
+        c.add("svucs", before_host="not-found-host", Hostname="venkateswara.onion", User="other", Port=22,
+              ProxyCommand="nc -w 300 -x localhost:9050 %h %p")
+
+    with pytest.raises(ValueError):  # its there now
+        c.add("venkateswara")
+
+
 def test_save():
     import tempfile
     tc = os.path.join(tempfile.gettempdir(), "temp_ssh_config-4123")
@@ -169,26 +217,34 @@ def test_mapping_add_new_keys():
 
 
 def test_remove():
+    c = sshconf.read_ssh_config(test_config)
 
-    def lines(fn):
-        with open(fn, "r") as f:
-            return len(f.read().splitlines())
+    c.add("abc", forwardAgent="yes", unknownpropertylikethis="noway", Hostname="ssh.svuni.local",
+          user="mmccaa")
 
-    import tempfile
-    tc = os.path.join(tempfile.gettempdir(), "temp_ssh_config-123")
-    try:
-        c = sshconf.read_ssh_config(test_config)
+    c.add("def", forwardAgent="yes", unknownpropertylikethis="noway", Hostname="ssh.svuni.local",
+          user="mmccaa")
 
-        assert 14 == lines(test_config)
+    config1 = c.config()
+    hosts = list(c.hosts())
+    c.remove("abc")
+    config2 = c.config()
+    hosts2 = list(c.hosts())
+    c.remove("svu")
+    config3 = c.config()
+    hosts3 = list(c.hosts())
 
-        c.remove("svu")
-        c.write(tc)
+    assert "abc" in hosts
+    assert "abc" not in hosts2
 
-        assert 9 == lines(tc)
-        assert "# within-host-comment" not in c.config()
+    assert "Host abc" in config1
+    assert "Host abc" not in config2
+    assert "Host svu" not in config3
 
-    finally:
-        os.remove(tc)
+    assert hosts2 == ["*", "svu", "def"]  # test order
+    assert hosts3 == ["*", "def"]
+
+    assert "# within-host-comment" not in config3
 
 
 def test_read_duplicate_keys():