Add a flatten argument to ssh.libs

CHANGELOG.md

@@ -95,6 +95,7 @@ The table below shows which release corresponds to each branch, and what date th
 - [#2345][2345] Fix pwn constgrep when it matches a non-constant type
 - [#2338][2338] Fix: follow symlink for libs on ssh connection
 - [#2341][2341] Launch GDB correctly in iTerm on Mac
+- [#2268][2268] Add a `flatten` argument to `ssh.libs`
 
 [2242]: https://github.com/Gallopsled/pwntools/pull/2242
 [2277]: https://github.com/Gallopsled/pwntools/pull/2277
@@ -118,6 +119,7 @@ The table below shows which release corresponds to each branch, and what date th
 [2345]: https://github.com/Gallopsled/pwntools/pull/2345
 [2338]: https://github.com/Gallopsled/pwntools/pull/2338
 [2341]: https://github.com/Gallopsled/pwntools/pull/2341
+[2268]: https://github.com/Gallopsled/pwntools/pull/2268
 
 ## 4.12.0 (`beta`)
 

pwnlib/tubes/ssh.py

@@ -1810,20 +1810,42 @@ 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) and
+                ignore the remote directory structure. If there are duplicate
+                filenames, an error will be raised.
+        """
 
         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
 
@@ -1834,7 +1856,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