diff --git a/common/configuration/puppet.yaml.tftpl b/common/configuration/puppet.yaml.tftpl
index 9a5a20b7..6dfb6afc 100644
--- a/common/configuration/puppet.yaml.tftpl
+++ b/common/configuration/puppet.yaml.tftpl
@@ -111,6 +111,7 @@ runcmd:
   - rm -rf /etc/puppetlabs/code/environments/production
   - git clone ${puppetenv_git} /etc/puppetlabs/code/environments/main
   - ln -s /etc/puppetlabs/code/environments/main /etc/puppetlabs/code/environments/production
+  - ln -s /etc/puppetlabs/code/environments/main /etc/puppetlabs/code/environments/image
   - |
     (
     cd /etc/puppetlabs/code/environments/production
@@ -137,6 +138,10 @@ runcmd:
 %{ for conf in puppet_conf ~}
   - ${jsonencode(["/opt/puppetlabs/bin/puppet", "config", "set", conf.key, conf.value, "--section", conf.section])}
 %{ endfor ~}
+%{ if contains(tags, "image") }
+  - /opt/puppetlabs/bin/puppet config --section agent set environment image
+  - /opt/puppetlabs/bin/puppet config set postrun_command /usr/sbin/prepare4image.sh
+%{ endif }
   - systemctl enable puppet
 # Remove all ifcfg configuration files that have no corresponding network interface in ip link show.
   - for i in /etc/sysconfig/network-scripts/ifcfg-*; do if ! ip link show | grep -q "$${i##*-}:"; then rm -f $i; fi; done