From 41681ee1cb0f1cd72a9b9d897cf133eea04b51b2 Mon Sep 17 00:00:00 2001
From: termanix <50464194+termanix@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:43:42 +0300
Subject: [PATCH] Fixed put-file

Signed-off-by: termanix <50464194+termanix@users.noreply.github.com>
---
 nxc/protocols/nfs.py | 37 +++++++++++++++++++++++++++----------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/nxc/protocols/nfs.py b/nxc/protocols/nfs.py
index d2f922f1..95ab7b08 100644
--- a/nxc/protocols/nfs.py
+++ b/nxc/protocols/nfs.py
@@ -1,5 +1,6 @@
 from nxc.connection import connection
 from nxc.logger import NXCAdapter
+from nxc.helpers.logger import highlight
 from pyNfsClient import Portmap, Mount, NFSv3, NFS_PROGRAM, NFS_V3, ACCESS3_READ, ACCESS3_MODIFY, ACCESS3_EXECUTE, NFSSTAT3
 import re
 import uuid
@@ -301,6 +302,11 @@ def put_file(self):
         remote_file_path = self.args.put_file[1]
         file_name = ""
 
+        # Check if local file is exist
+        if not os.path.isfile(local_file_path):
+            self.logger.fail(f"{local_file_path} does not exist.")
+            return
+
         # Do a bit of smart handling for the file paths
         file_name = local_file_path.split("/")[-1] if "/" in local_file_path else local_file_path
         if not remote_file_path.endswith("/"):
@@ -313,14 +319,18 @@ def put_file(self):
             self.nfs3 = NFSv3(self.host, nfs_port, self.args.nfs_timeout, self.auth)
             self.nfs3.connect()
 
-            try:
-                # Mount the NFS share to create the file
-                mnt_info = self.mount.mnt(remote_file_path, self.auth)
-                dir_handle = mnt_info["mountinfo"]["fhandle"]
-                # Update the UID from the directory
-                attrs = self.nfs3.getattr(dir_handle, auth=self.auth)
-                self.auth["uid"] = attrs["attributes"]["uid"]
+            # Mount the NFS share to create the file
+            mnt_info = self.mount.mnt(remote_file_path, self.auth)
+            dir_handle = mnt_info["mountinfo"]["fhandle"]
+            # Update the UID from the directory
+            attrs = self.nfs3.getattr(dir_handle, auth=self.auth)
+            self.auth["uid"] = attrs["attributes"]["uid"]
+            
+            # Checking if file_name is already exist on remote file path
+            lookup_response = self.nfs3.lookup(dir_handle, file_name, auth=self.auth)
 
+            # If success, file_name does not exist on remote machine. Else, trying to overwrite it.
+            if lookup_response["resok"] is None:
                 # Create file
                 self.logger.display(f"Trying to create {remote_file_path}{file_name}")
                 res = self.nfs3.create(dir_handle, file_name, create_mode=1, mode=0o777, auth=self.auth)
@@ -329,9 +339,16 @@ def put_file(self):
                 else:
                     file_handle = res["resok"]["obj"]["handle"]["data"]
                 self.logger.success(f"{file_name} successfully created")
-            except Exception as e:
-                self.logger.fail(f"{file_name} was not created: {e}")
-                return
+            else:
+                # Asking to user overwriting.
+                ans = input(highlight(f"[!] {file_name} is already exist on {remote_file_path}. Do you want to overwrite it? [Y/n] ", "red"))
+                if ans.lower() in ["y", "yes", ""]:
+                    # Overwriting
+                    self.logger.display(f"{file_name} is already exist on {remote_file_path}. Trying to overwrite it...")
+                    file_handle = lookup_response["resok"]["object"]["data"]
+                else:
+                    self.logger.fail(f"Uploading was not successful. The {file_name} is exist on {remote_file_path}")
+                    return
 
             try:
                 with open(local_file_path, "rb") as file: