Add a flatten argument to ssh.libs

CHANGELOG.md

@@ -89,6 +89,7 @@ The table below shows which release corresponds to each branch, and what date th
 - [#2212][2212] Add `--libc libc.so` argument to `pwn template` command
 - [#2257][2257] Allow creation of custom templates for `pwn template` command
 - [#2225][2225] Allow empty argv in ssh.process()
+- [#2268][2268] Add a `flatten` argument to `ssh.libs`
 
 [2202]: https://github.com/Gallopsled/pwntools/pull/2202
 [2117]: https://github.com/Gallopsled/pwntools/pull/2117
@@ -97,6 +98,7 @@ The table below shows which release corresponds to each branch, and what date th
 [2212]: https://github.com/Gallopsled/pwntools/pull/2212
 [2257]: https://github.com/Gallopsled/pwntools/pull/2257
 [2225]: https://github.com/Gallopsled/pwntools/pull/2225
+[2268]: https://github.com/Gallopsled/pwntools/pull/2268
 
 ## 4.11.1 (`stable`)
 

pwnlib/tubes/ssh.py

@@ -1804,20 +1804,53 @@ def unlink(self, file):
 
         return self.sftp.unlink(file)
 
-    def libs(self, remote, directory = None):
+    def libs(self, remote, directory = None, flatten = False):
         """Downloads the libraries referred to by a file.
 
         This is done by running ldd on the remote server, parsing the output
         and downloading the relevant files.
 
         The directory argument specified where to download the files. This defaults
-        to './$HOSTNAME' where $HOSTNAME is the hostname of the remote server."""
+        to './$HOSTNAME' where $HOSTNAME is the hostname of the remote server.
+
+        Arguments:
+            remote(str): Remote file path
+            directory(str): Output directory
+            flatten(bool): Flatten the file tree if True (defaults to False). If
+                there are duplicate filenames, fallback to unflattened structure.
+
+        Examples:
+            >>> s =  ssh(host='example.pwnme')
+            >>> s.cwd = f'{pwnlib.data.elf.path}/ssh_libs'
+            >>> s.libs("duplicate", "out_duplicate") # doctest: +ELLIPSIS
+            {'.../b/lib.so': ..., '.../a/lib.so': ..., '.../duplicate': ...}
+            >>> s.libs("no_duplicate", "out_noduplicate_flatten", flatten = True) # doctest: +ELLIPSIS
+            {'.../out_noduplicate_flatten/lib.so': ..., '.../out_noduplicate_flatten/lib2.so': ..., '.../out_noduplicate_flatten/no_duplicate': ...}
+            >>> s.libs("duplicate", "out_duplicate_flatten", flatten = True) # doctest: +ELLIPSIS
+            Traceback (most recent call last):
+            ...
+            PwnlibException: Duplicate lib name: ...
+        """
 
         libs = self._libs_remote(remote)
 
         remote = packing._decode(self.readlink('-f',remote).strip())
         libs[remote] = 0
 
+        if flatten:
+            basenames = dict()
+
+            # If there is a duplicate switch to unflattened download
+            for lib in libs:
+                name = os.path.basename(lib)
+
+                if name in basenames.values():
+                    duplicate = [key for key, value in basenames.items() if
+                                 value == name][0]
+                    self.error('Duplicate lib name: %r / %4r' % (lib, duplicate))
+
+                basenames[lib] = name
+
         if directory is None:
             directory = self.host
 
@@ -1828,7 +1861,8 @@ def libs(self, remote, directory = None):
         seen = set()
 
         for lib, addr in libs.items():
-            local = os.path.realpath(os.path.join(directory, '.' + os.path.sep + lib))
+            local = os.path.realpath(os.path.join(directory, '.' + os.path.sep \
+                    + (basenames[lib] if flatten else lib)))
             if not local.startswith(directory):
                 self.warning('This seems fishy: %r' % lib)
                 continue