From 6b7cb7bd38348dbe4a56385356df6ed97e81c319 Mon Sep 17 00:00:00 2001
From: Linus <23507341+Linus789@users.noreply.github.com>
Date: Fri, 10 Nov 2023 01:05:36 +0000
Subject: [PATCH] fix(Remove screenshot restriction): Improve reliability
 (#2938)

Co-authored-by: oSumAtrIX <johan.melkonyan1@web.de>
---
 .../RemoveScreenshotRestrictionPatch.kt       | 50 +++++++++++++++++++
 .../app/revanced/util/patch/MethodCall.kt     |  1 +
 2 files changed, 51 insertions(+)

diff --git a/src/main/kotlin/app/revanced/patches/all/screenshot/removerestriction/RemoveScreenshotRestrictionPatch.kt b/src/main/kotlin/app/revanced/patches/all/screenshot/removerestriction/RemoveScreenshotRestrictionPatch.kt
index ce0db09da..a16187681 100644
--- a/src/main/kotlin/app/revanced/patches/all/screenshot/removerestriction/RemoveScreenshotRestrictionPatch.kt
+++ b/src/main/kotlin/app/revanced/patches/all/screenshot/removerestriction/RemoveScreenshotRestrictionPatch.kt
@@ -1,14 +1,19 @@
 package app.revanced.patches.all.screenshot.removerestriction
 
+import app.revanced.patcher.data.BytecodeContext
+import app.revanced.patcher.extensions.InstructionExtensions.addInstructions
 import app.revanced.patcher.patch.annotation.Patch
 import app.revanced.patcher.util.proxy.mutableTypes.MutableMethod
 import app.revanced.util.patch.AbstractTransformInstructionsPatch
 import app.revanced.util.patch.IMethodCall
 import app.revanced.util.patch.Instruction35cInfo
 import app.revanced.util.patch.filterMapInstruction35c
+import com.android.tools.smali.dexlib2.Opcode
 import com.android.tools.smali.dexlib2.iface.ClassDef
 import com.android.tools.smali.dexlib2.iface.Method
 import com.android.tools.smali.dexlib2.iface.instruction.Instruction
+import com.android.tools.smali.dexlib2.iface.instruction.formats.Instruction22c
+import com.android.tools.smali.dexlib2.iface.reference.FieldReference
 
 @Patch(
     name = "Remove screenshot restriction",
@@ -22,6 +27,11 @@ object RemoveScreenshotRestrictionPatch : AbstractTransformInstructionsPatch<Ins
             "Lapp/revanced/all/screenshot/removerestriction/RemoveScreenshotRestrictionPatch"
     private const val INTEGRATIONS_CLASS_DESCRIPTOR = "$INTEGRATIONS_CLASS_DESCRIPTOR_PREFIX;"
 
+    override fun execute(context: BytecodeContext) {
+        super.execute(context)
+        ModifyLayoutParamsFlags().execute(context)
+    }
+
     override fun filterMap(
         classDef: ClassDef,
         method: Method,
@@ -46,6 +56,12 @@ object RemoveScreenshotRestrictionPatch : AbstractTransformInstructionsPatch<Ins
         override val methodParams: Array<String>,
         override val returnType: String
     ): IMethodCall {
+        AddFlags(
+            "Landroid/view/Window;",
+            "addFlags",
+            arrayOf("I"),
+            "V",
+        ),
         SetFlags(
             "Landroid/view/Window;",
             "setFlags",
@@ -54,3 +70,37 @@ object RemoveScreenshotRestrictionPatch : AbstractTransformInstructionsPatch<Ins
         );
     }
 }
+
+private class ModifyLayoutParamsFlags : AbstractTransformInstructionsPatch<Pair<Instruction22c, Int>>() {
+    override fun filterMap(
+        classDef: ClassDef,
+        method: Method,
+        instruction: Instruction,
+        instructionIndex: Int
+    ): Pair<Instruction22c, Int>? {
+        if (instruction.opcode != Opcode.IPUT) {
+            return null
+        }
+
+        val instruction22c = instruction as Instruction22c
+        val fieldReference = instruction22c.reference as FieldReference
+
+        if (fieldReference.definingClass != "Landroid/view/WindowManager\$LayoutParams;"
+            || fieldReference.name != "flags"
+            || fieldReference.type != "I") {
+            return null
+        }
+
+        return Pair(instruction22c, instructionIndex)
+    }
+
+    override fun transform(mutableMethod: MutableMethod, entry: Pair<Instruction22c, Int>) {
+        val (instruction, index) = entry
+        val register = instruction.registerA
+
+        mutableMethod.addInstructions(
+            index,
+            "and-int/lit16 v$register, v$register, -0x2001"
+        )
+    }
+}
diff --git a/src/main/kotlin/app/revanced/util/patch/MethodCall.kt b/src/main/kotlin/app/revanced/util/patch/MethodCall.kt
index 48a38ad07..1533c1ca7 100644
--- a/src/main/kotlin/app/revanced/util/patch/MethodCall.kt
+++ b/src/main/kotlin/app/revanced/util/patch/MethodCall.kt
@@ -67,6 +67,7 @@ inline fun <reified E> fromMethodReference(methodReference: MethodReference)
     search.definedClassName == methodReference.definingClass
             && search.methodName == methodReference.name
             && methodReference.parameterTypes.toTypedArray().contentEquals(search.methodParams)
+            && search.returnType == methodReference.returnType
 }
 
 inline fun <reified E> filterMapInstruction35c(