diff --git a/bin/waagent b/bin/waagent
index 4039e037d2..60cdf56c56 100755
--- a/bin/waagent
+++ b/bin/waagent
@@ -24,9 +24,13 @@
 #
 
 import os
-import imp
 import sys
 
+if sys.version_info[0] == 2:
+    import imp
+else:
+    import importlib
+
 if __name__ == '__main__' :
     import azurelinuxagent.agent as agent
     """
diff --git a/bin/waagent2.0 b/bin/waagent2.0
index 25aa0ce0ff..347326776c 100644
--- a/bin/waagent2.0
+++ b/bin/waagent2.0
@@ -1744,7 +1744,11 @@ bsd_activate_resource_disk_txt="""\
 
 import os
 import sys
-import imp
+
+if sys.version_info[0] == 2:
+    import imp
+else:
+    import importlib
 
 # waagent has no '.py' therefore create waagent module import manually.
 __name__='setupmain' #prevent waagent.__main__ from executing