From 6500ff68f328aa6c009603ddac144a57b0db4685 Mon Sep 17 00:00:00 2001
From: Erik Costlow <erikcostlow@outlook.com>
Date: Fri, 17 Dec 2021 10:01:21 -0500
Subject: [PATCH] Filename detection

---
 src/main/java/com/contrastsecurity/App.java | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/contrastsecurity/App.java b/src/main/java/com/contrastsecurity/App.java
index 6f8b21e..7deebdc 100644
--- a/src/main/java/com/contrastsecurity/App.java
+++ b/src/main/java/com/contrastsecurity/App.java
@@ -1,18 +1,15 @@
 package com.contrastsecurity;
 
 import java.io.File;
-import java.io.IOException;
-import java.net.URISyntaxException;
 import java.util.List;
-import net.bytebuddy.agent.ByteBuddyAgent;
 
-import com.sun.tools.attach.AgentInitializationException;
-import com.sun.tools.attach.AgentLoadException;
-import com.sun.tools.attach.AttachNotSupportedException;
 import com.sun.tools.attach.VirtualMachine;
 import com.sun.tools.attach.VirtualMachineDescriptor;
 
+import net.bytebuddy.agent.ByteBuddyAgent;
+
 public class App {
+
     public static void main(String[] args){
         System.out.println("SafeLog4j by Contrast Security");
         System.out.println( "https://contrastsecurity.com" );
@@ -21,8 +18,13 @@ public static void main(String[] args){
             try{
                 String pid = args[0];
                 String options = args.length>=2 ? args[1] : null;
-                String agentFilePath = "safelog4j-1.0.1.jar";
-                File agentFile = new File(agentFilePath);
+
+                String filename = App.class.getProtectionDomain()
+                    .getCodeSource()
+                    .getLocation()
+                    .toURI()
+                    .getPath();
+                File agentFile = new File(filename);
                 ByteBuddyAgent.attach(agentFile.getAbsoluteFile(), pid, options);
                 System.out.println("Attached to target jvm and loaded agent successfully");
             }catch(Exception e){