From 679ea19d6c6bafe6151c86fcaaec682c1ca54d4a Mon Sep 17 00:00:00 2001
From: Dustin Oprea <myselfasunder@gmail.com>
Date: Wed, 18 Apr 2018 07:54:47 -0400
Subject: [PATCH] ifd_enumerate. The IFD tree and content can now be collected.

- Collected into a static structure in addition to scanned (which is
- only a visitor pattern).
- Test still has be finished.
---
 exif.go          |  56 +++++++++++++++++
 exif_test.go     |  36 +++++++++++
 ifd_enumerate.go | 159 ++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 249 insertions(+), 2 deletions(-)

diff --git a/exif.go b/exif.go
index 5d55142..cc2a0ea 100644
--- a/exif.go
+++ b/exif.go
@@ -1,9 +1,11 @@
 package exif
 
 import (
+    "os"
     "errors"
     "bytes"
 
+    "io/ioutil"
     "encoding/binary"
 
     "github.com/dsoprea/go-logging"
@@ -34,6 +36,42 @@ func (e *Exif) IsExif(data []byte) (ok bool) {
     return false
 }
 
+func (e *Exif) SearchAndExtractExif(filepath string) (rawExif []byte, err error) {
+    defer func() {
+        if state := recover(); state != nil {
+            err := log.Wrap(state.(error))
+            log.Panic(err)
+        }
+    }()
+
+    // Open the file.
+
+    f, err := os.Open(filepath)
+    log.PanicIf(err)
+
+    defer f.Close()
+
+    data, err := ioutil.ReadAll(f)
+    log.PanicIf(err)
+
+    // Search for the beginning of the EXIF information. The EXIF is near the
+    // beginning of our/most JPEGs, so this has a very low cost.
+
+    foundAt := -1
+    for i := 0; i < len(data); i++ {
+        if e.IsExif(data[i:i + 6]) == true {
+            foundAt = i
+            break
+        }
+    }
+
+    if foundAt == -1 {
+        log.Panicf("EXIF start not found")
+    }
+
+    return data[foundAt:], nil
+}
+
 
 type ExifHeader struct {
     ByteOrder binary.ByteOrder
@@ -103,3 +141,21 @@ func (e *Exif) Visit(data []byte, visitor TagVisitor) (err error) {
 
     return nil
 }
+
+func (e *Exif) Collect(data []byte) (rootIfd *Ifd, tree map[int]*Ifd, ifds []*Ifd, err error) {
+    defer func() {
+        if state := recover(); state != nil {
+            err = log.Wrap(state.(error))
+        }
+    }()
+
+    eh, err := e.ParseExifHeader(data)
+    log.PanicIf(err)
+
+    ie := NewIfdEnumerate(data, eh.ByteOrder)
+
+    rootIfd, tree, ifds, err = ie.Collect(eh.FirstIfdOffset)
+    log.PanicIf(err)
+
+    return rootIfd, tree, ifds, nil
+}
diff --git a/exif_test.go b/exif_test.go
index d5432f7..b82ae62 100644
--- a/exif_test.go
+++ b/exif_test.go
@@ -187,6 +187,42 @@ func TestVisit(t *testing.T) {
     }
 }
 
+func TestCollect(t *testing.T) {
+    defer func() {
+        if state := recover(); state != nil {
+            err := log.Wrap(state.(error))
+            log.PrintErrorf(err, "Exif failure.")
+        }
+    }()
+
+    e := NewExif()
+
+    filepath := path.Join(assetsPath, "NDM_8901.jpg")
+
+    rawExif, err := e.SearchAndExtractExif(filepath)
+    log.PanicIf(err)
+
+    rootIfd, tree, ifds, err := e.Collect(rawExif)
+    log.PanicIf(err)
+
+    rootIfd = rootIfd
+    tree = tree
+    ifds = ifds
+
+
+// TODO(dustin): !! Finish test.
+
+
+    // rootIfd.PrintTree()
+
+    // expected := []string {
+    // }
+
+    // if reflect.DeepEqual(tags, expected) == false {
+    //     t.Fatalf("tags not correct:\n%v", tags)
+    // }
+}
+
 func init() {
     goPath := os.Getenv("GOPATH")
     if goPath == "" {
diff --git a/ifd_enumerate.go b/ifd_enumerate.go
index d5d1d04..6ad9b60 100644
--- a/ifd_enumerate.go
+++ b/ifd_enumerate.go
@@ -2,6 +2,8 @@ package exif
 
 import (
     "bytes"
+    "fmt"
+    "strings"
 
     "encoding/binary"
 
@@ -132,7 +134,7 @@ type IfdTagEntry struct {
     UnitCount uint32
     ValueOffset uint32
     RawValueOffset []byte
-    IsIfd bool
+    IfdName string
 }
 
 
@@ -207,7 +209,7 @@ func (ie *IfdEnumerate) ParseIfd(ifdName string, ifdIndex int, ifdOffset uint32,
 
         childIfdName, isIfd := IsIfdTag(tagId)
         if isIfd == true {
-            tag.IsIfd = true
+            tag.IfdName = childIfdName
 
             if doDescend == true {
                 ifdLogger.Debugf(nil, "Descending to IFD [%s].", childIfdName)
@@ -249,3 +251,156 @@ func (ie *IfdEnumerate) Scan(ifdName string, ifdOffset uint32, visitor TagVisito
 
     return nil
 }
+
+
+type Ifd struct {
+    Id int
+    ParentIfd *Ifd
+    Name string
+    Index int
+    Offset uint32
+    Entries []IfdTagEntry
+    Children []*Ifd
+    NextIfdOffset uint32
+    NextIfd *Ifd
+}
+
+func (ifd Ifd) String() string {
+    parentOffset := uint32(0)
+    if ifd.ParentIfd != nil {
+        parentOffset = ifd.ParentIfd.Offset
+    }
+
+    return fmt.Sprintf("IFD<ID=(%d) N=[%s] IDX=(%d) OFF=(0x%04x) COUNT=(%d) CHILDREN=(%d) PARENT=(0x%04x) NEXT-IFD=(0x%04x)", ifd.Id, ifd.Name, ifd.Index, ifd.Offset, len(ifd.Entries), len(ifd.Children), parentOffset, ifd.NextIfdOffset)
+}
+
+func (ifd Ifd) printNode(level int, nextLink bool) {
+    indent := strings.Repeat(" ", level * 2)
+
+    prefix := " "
+    if nextLink {
+        prefix = ">"
+    }
+
+    fmt.Printf("%s%s%s\n", indent, prefix, ifd)
+
+    for _, childIfd := range ifd.Children {
+        childIfd.printNode(level + 1, false)
+    }
+
+    if ifd.NextIfd != nil {
+        ifd.NextIfd.printNode(level, true)
+    }
+}
+
+func (ifd Ifd) PrintTree() {
+    ifd.printNode(0, false)
+}
+
+
+type QueuedIfd struct {
+    Name string
+    Index int
+    Offset uint32
+    Parent *Ifd
+}
+
+// Scan enumerates the different EXIF blocks (called IFDs).
+func (ie *IfdEnumerate) Collect(rootIfdOffset uint32) (rootIfd *Ifd, tree map[int]*Ifd, ifds []*Ifd, err error) {
+    defer func() {
+        if state := recover(); state != nil {
+            err = log.Wrap(state.(error))
+        }
+    }()
+
+    tree = make(map[int]*Ifd)
+    ifds = make([]*Ifd, 0)
+
+    queue := []QueuedIfd {
+        {
+            Name: IfdStandard,
+            Index: 0,
+            Offset: rootIfdOffset,
+        },
+    }
+
+    edges := make(map[uint32]*Ifd)
+
+    for {
+        if len(queue) == 0 {
+            break
+        }
+
+        name := queue[0].Name
+        index := queue[0].Index
+        offset := queue[0].Offset
+        parentIfd := queue[0].Parent
+
+        queue = queue[1:]
+
+        nextIfdOffset, entries, err := ie.ParseIfd(name, index, offset, nil, false)
+        log.PanicIf(err)
+
+        id := len(ifds)
+
+        ifd := Ifd{
+            Id: id,
+            ParentIfd: parentIfd,
+            Name: name,
+            Index: index,
+            Offset: offset,
+            Entries: entries,
+            Children: make([]*Ifd, 0),
+            NextIfdOffset: nextIfdOffset,
+        }
+
+        // Add ourselves to a big list of IFDs.
+        ifds = append(ifds, &ifd)
+
+        // Install ourselves into a lookup table.
+        tree[id] = &ifd
+
+        // Add a link from the previous IFD in the chain to us.
+        if previousIfd, found := edges[offset]; found == true {
+            previousIfd.NextIfd = &ifd
+        }
+
+        // Attach as a child to our parent (where we appeared as a tag in
+        // that IFD).
+        if parentIfd != nil {
+            parentIfd.Children = append(parentIfd.Children, &ifd)
+        }
+
+        // Determine if any of our entries is a child IFD and queue it.
+        for _, entry := range entries {
+            if entry.IfdName == "" {
+                continue
+            }
+
+            qi := QueuedIfd {
+                Name: entry.IfdName,
+                Index: 0,
+                Offset: entry.ValueOffset,
+                Parent: &ifd,
+            }
+
+            queue = append(queue, qi)
+        }
+
+        // If there's another IFD in the chain.
+        if nextIfdOffset != 0 {
+            // Allow the next link to know what the previous link was.
+            edges[nextIfdOffset] = &ifd
+
+            qi := QueuedIfd {
+                Name: IfdStandard,
+                Index: index + 1,
+                Offset: nextIfdOffset,
+            }
+
+            queue = append(queue, qi)
+        }
+    }
+
+    return tree[0], tree, ifds, nil
+}